You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "Kimahriman (via GitHub)" <gi...@apache.org> on 2024/09/09 16:34:43 UTC

[PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Kimahriman opened a new pull request, #48038:
URL: https://github.com/apache/spark/pull/48038

   <!--
   Thanks for sending a pull request!  Here are some tips for you:
     1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
     2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
     3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
     4. Be sure to keep the PR description updated to reflect all changes.
     5. Please write your PR title to summarize what this PR proposes.
     6. If possible, provide a concise example to reproduce the issue for a faster review.
     7. If you want to add a new configuration, please read the guideline first for naming configurations in
        'core/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
     8. If you want to add or modify an error type or message, please read the guideline first in
        'common/utils/src/main/resources/error/README.md'.
   -->
   
   ### What changes were proposed in this pull request?
   <!--
   Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue. 
   If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
     1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
     2. If you fix some SQL features, you can provide some references of other DBMSes.
     3. If there is design documentation, please add the link.
     4. If there is a discussion in the mailing list, please add the link.
   -->
   Add the option to `applyInArrow` to return an iterator of RecordBatches instead of a single pyarrow Table.
   
   ### Why are the changes needed?
   <!--
   Please clarify why the changes are needed. For instance,
     1. If you propose a new API, clarify the use case for a new API.
     2. If you fix a bug, you can clarify why it is a bug.
   -->
   Being limited to returning a single Table requires collecting all results in memory for a single batch. This can require excessive memory for certain edge cases that might require returning a larger number of rows.
   
   ### Does this PR introduce _any_ user-facing change?
   <!--
   Note that it means *any* user-facing change including all aspects such as the documentation fix.
   If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
   If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
   If no, write 'No'.
   -->
   Yes, an additional way to return data from applyInArrow
   
   ### How was this patch tested?
   <!--
   If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
   If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
   If tests were not added, please describe why they were not added and/or why it was difficult to add.
   If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
   -->
   Updated existing UTs to test both returning a Table and iterator of RecordBatch
   
   ### Was this patch authored or co-authored using generative AI tooling?
   <!--
   If generative AI tooling has been used in the process of authoring this patch, please include the
   phrase: 'Generated-by: ' followed by the name of the tool and its version.
   If no, write 'No'.
   Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
   -->
   No


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2341604597

   I updated the grouped API to be both Table -> Table or Iterator -> Iterator. I guess I will do the same for cogrouped, maybe some day they could actually be iterated over independently so might as well keep the APIs consistent


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "EnricoMi (via GitHub)" <gi...@apache.org>.
EnricoMi commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1775246226


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -803,27 +806,30 @@ def applyInArrow(
         Applies a function to each cogroup using Arrow and returns the result
         as a `DataFrame`.
 
-        The function should take two `pyarrow.Table`\\s and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and the two `pyarrow.Table`\\s.
-        For each side of the cogroup, all columns are passed together as a
-        `pyarrow.Table` to the user-function and the returned `pyarrow.Table` are combined as
-        a :class:`DataFrame`.
+        The function can take one of two forms: It can take two `pyarrow.Table`\\s and return a
+        `pyarrow.Table`, or it can take two iterators of `pyarrow.RecordBatch` and yield
+        `pyarrow.RecordBatch`. Alternatively, each form can take a tuple of `pyarrow.Scalar`
+        as the first argument in addition to the input type above. For each cogroup, all columns
+        are passed together in the `pyarrow.Table` or `pyarrow.RecordBatch`, and the returned
+        `pyarrow.Table` or iterator of `pyarrow.RecordBatch` are combined as a :class:`DataFrame`.
 
         The `schema` should be a :class:`StructType` describing the schema of the returned

Review Comment:
   Maybe a follow-up PR to keep this one tidy?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2405924462

   Gentle ping @zhengruifeng @HyukjinKwon 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "EnricoMi (via GitHub)" <gi...@apache.org>.
EnricoMi commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2341595186

   > Both left and right have to be serialized over the same channel
   
   That is my understanding as well. The two sides would have to be transferred over two channels so you can independently iterate over the iterators. Currently, any JVM <-> Python communication serializes over a single channel.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1775511804


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -803,27 +806,30 @@ def applyInArrow(
         Applies a function to each cogroup using Arrow and returns the result
         as a `DataFrame`.
 
-        The function should take two `pyarrow.Table`\\s and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and the two `pyarrow.Table`\\s.
-        For each side of the cogroup, all columns are passed together as a
-        `pyarrow.Table` to the user-function and the returned `pyarrow.Table` are combined as
-        a :class:`DataFrame`.
+        The function can take one of two forms: It can take two `pyarrow.Table`\\s and return a
+        `pyarrow.Table`, or it can take two iterators of `pyarrow.RecordBatch` and yield
+        `pyarrow.RecordBatch`. Alternatively, each form can take a tuple of `pyarrow.Scalar`
+        as the first argument in addition to the input type above. For each cogroup, all columns
+        are passed together in the `pyarrow.Table` or `pyarrow.RecordBatch`, and the returned
+        `pyarrow.Table` or iterator of `pyarrow.RecordBatch` are combined as a :class:`DataFrame`.
 
         The `schema` should be a :class:`StructType` describing the schema of the returned

Review Comment:
   Yeah I can look into that in my follow on that actually does the batching on the JVM input side (unless others just want me to merge that into this branch)



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1751352741


##########
python/pyspark/sql/pandas/_typing/__init__.pyi:
##########
@@ -347,10 +348,16 @@ PandasCogroupedMapFunction = Union[
 ArrowGroupedMapFunction = Union[
     Callable[[pyarrow.Table], pyarrow.Table],
     Callable[[Tuple[pyarrow.Scalar, ...], pyarrow.Table], pyarrow.Table],
+    Callable[[pyarrow.Table], Iterator[pyarrow.RecordBatch]],
+    Callable[[Tuple[pyarrow.Scalar, ...], pyarrow.Table], Iterator[pyarrow.RecordBatch]],

Review Comment:
   +1, I also prefer `iterator -> iterator`



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2500828857

   are we still using `FlatMapGroupsInArrowExec` which need to have enough memory to
   hold the largest group?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "EnricoMi (via GitHub)" <gi...@apache.org>.
EnricoMi commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2338882363

   I'd vote for having iterator of RecordBatches be the only option for input and output. Given we are at a major release, we can <strike>break</strike> evolve the API. You also point out rightly that this would be a big chunk of work and there is not time until the release. But Spark devs have to comment on their appetite to touch the Python API.
   
   As you said, this wouldn't work for cogroups, leading to some inconsistency in the API.
   
   What you propose is what exists for the Scala / Java API (ignore the sort bit): iterating the groups to implement memory-efficient stream algorithms: https://github.com/apache/spark/pull/39640 Would be fair to add this to Python.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1752718645


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -538,26 +538,28 @@ def applyInArrow(
         as a `DataFrame`.
 
         The function should take a `pyarrow.Table` and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and a `pyarrow.Table`.
-        For each group, all columns are passed together as a `pyarrow.Table`
-        to the user-function and the returned `pyarrow.Table` are combined as a
-        :class:`DataFrame`.
+        `pyarrow.Table` or an iterator of `pyarrow.RecordBatch. Alternatively,
+        the user can pass a function that takes a tuple of `pyarrow.Scalar` grouping
+        key(s) and a `pyarrow.Table`. For each group, all columns are passed together
+        as a`pyarrow.Table` to the user-function and the returned `pyarrow.Table` or iterator

Review Comment:
   Fully rewrote the doc



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1774919840


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -803,27 +806,30 @@ def applyInArrow(
         Applies a function to each cogroup using Arrow and returns the result
         as a `DataFrame`.
 
-        The function should take two `pyarrow.Table`\\s and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and the two `pyarrow.Table`\\s.
-        For each side of the cogroup, all columns are passed together as a
-        `pyarrow.Table` to the user-function and the returned `pyarrow.Table` are combined as
-        a :class:`DataFrame`.
+        The function can take one of two forms: It can take two `pyarrow.Table`\\s and return a
+        `pyarrow.Table`, or it can take two iterators of `pyarrow.RecordBatch` and yield
+        `pyarrow.RecordBatch`. Alternatively, each form can take a tuple of `pyarrow.Scalar`
+        as the first argument in addition to the input type above. For each cogroup, all columns
+        are passed together in the `pyarrow.Table` or `pyarrow.RecordBatch`, and the returned
+        `pyarrow.Table` or iterator of `pyarrow.RecordBatch` are combined as a :class:`DataFrame`.
 
         The `schema` should be a :class:`StructType` describing the schema of the returned

Review Comment:
   That's not currently true with how it's implemented now, but I can look into making that happen (and add the doc). Currently the serializer materializes both left and right fully into memory.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2342406217

   > > because I guess taking Table as input won't fix this issue
   > 
   > I can make the API change for cogrouped as well, but I don't think it's really going to be possible to to fix the loading all data into memory issue. Both left and right have to be serialized over the same channel, so you can't partially load the left and the right at the same time, you need to fully load the left at least first AFAIK. It would just be cosmetic but always an iterator of a single RecordBatch already loaded into memory
   
   good question, I didn't realize this issue before.
   I am wondering is it possible to zip the two iterators in the JVM side and then unzip it in Python side?
   
   cc @HyukjinKwon @ueshin @xinrong-meng  do you have some ideas on the cogrouped API?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2341826616

   I updated the cogroup API as well. I think that is all I am trying to tackle with this PR. The APIs are updated to support iterator -> iterator methods, and this solves the current issue of being able to yield multiple batches from the UDF. It also is prepared for taking multiple batches as input once the serialization supports it. As a follow on I can investigate actually implementing the batching logic on the serialization side, unless anyone feels strongly this PR should wait for and include that functionality as well.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2339487565

   +1 for `Iterator[RecordBatch] -> Iterator[RecordBatch]` to handle large group, which is consistent with [mapInArrow](https://spark.apache.org/docs/4.0.0-preview1/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.mapInArrow.html?highlight=mapinarrow#pyspark.sql.DataFrame.mapInArrow)
   
   and I suggest we still keep existing `Table -> Table` interface, it should be able to support both.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1774947965


##########
python/pyspark/worker.py:
##########
@@ -333,17 +336,38 @@ def wrap_cogrouped_map_arrow_udf(f, return_type, argspec, runner_conf):
             (col.name, to_arrow_type(col.dataType)) for col in return_type.fields
         ]
 
-    def wrapped(left_key_table, left_value_table, right_key_table, right_value_table):
-        if len(argspec.args) == 2:
-            result = f(left_value_table, right_value_table)
-        elif len(argspec.args) == 3:
-            key_table = left_key_table if left_key_table.num_rows > 0 else right_key_table
-            key = tuple(c[0] for c in key_table.columns)
-            result = f(key, left_value_table, right_value_table)
-
-        verify_arrow_result(result, _assign_cols_by_name, expected_cols_and_types)
+    def wrapped(left_key_batch, left_value_batches, right_key_batch, right_value_batches):
+        if is_generator:
+            if len(argspec.args) == 2:
+                result = f(left_value_batches, right_value_batches)
+            elif len(argspec.args) == 3:
+                key_batch = left_key_batch if left_key_batch.num_rows > 0 else right_key_batch
+                key = tuple(c[0] for c in key_batch.columns)
+                result = f(key, left_value_batches, right_value_batches)
+
+            def verify_element(batch):
+                verify_arrow_batch(batch, _assign_cols_by_name, expected_cols_and_types)
+                return batch
+
+            yield from map(verify_element, result)
+            # Make sure both iterators are fully consumed
+            for _ in left_value_batches:
+                pass
+            for _ in right_value_batches:
+                pass
 
-        return result.to_batches()
+        else:

Review Comment:
   Good idea, will update



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2500843684

   > Why there is no jvm side execution change? Is it still using `FlatMapGroupsInArrowExec` which need to have enough memory to hold the largest group?
   
   I have a follow up ready that implements the input side batching, was trying to keep this PR more manageable 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2506184445

   I made a new PR with all the changes combined: https://github.com/apache/spark/pull/49005
   
   All the connect support mostly happens for free. I did not try to touch the pandas APIs, I think that could be a separate followup.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1779744350


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -803,27 +806,30 @@ def applyInArrow(
         Applies a function to each cogroup using Arrow and returns the result
         as a `DataFrame`.
 
-        The function should take two `pyarrow.Table`\\s and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and the two `pyarrow.Table`\\s.
-        For each side of the cogroup, all columns are passed together as a
-        `pyarrow.Table` to the user-function and the returned `pyarrow.Table` are combined as
-        a :class:`DataFrame`.
+        The function can take one of two forms: It can take two `pyarrow.Table`\\s and return a
+        `pyarrow.Table`, or it can take two iterators of `pyarrow.RecordBatch` and yield
+        `pyarrow.RecordBatch`. Alternatively, each form can take a tuple of `pyarrow.Scalar`
+        as the first argument in addition to the input type above. For each cogroup, all columns
+        are passed together in the `pyarrow.Table` or `pyarrow.RecordBatch`, and the returned
+        `pyarrow.Table` or iterator of `pyarrow.RecordBatch` are combined as a :class:`DataFrame`.
 
         The `schema` should be a :class:`StructType` describing the schema of the returned

Review Comment:
   @zhengruifeng this is a potential way to address your comment above. Just load the left side fully into memory, but iteratively load the right side, so you at least don't need both in memory at once.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2338583728

   @EnricoMi @HyukjinKwon 
   
   I made this backward compatible with the current implementation. Some alternative things that could be done if this is still considered an unreleased feature (not sure if a preview release counts) would be:
   - Make an iterator of `RecordBatch` be the only return type. AFAIK there's nothing to gain using a single Table over RecordBatches, since you can easily turn a Table into batches, and that's all the worker does immediately with the result
   - Update the input to also be an iterator of `RecordBatch`. I considered this here but it would require much more updating to the serialization process, as a lot of things assume a single batch per group. It should be possible to do but will require more work, and it wouldn't apply to cogroups, since we need to serialize the left then the right.
   
   Happy to look more into either alternative if you think it makes sense.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2339520424

   i'm fine with this.


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "HyukjinKwon (via GitHub)" <gi...@apache.org>.
HyukjinKwon commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1751190917


##########
python/pyspark/sql/pandas/_typing/__init__.pyi:
##########
@@ -347,10 +348,16 @@ PandasCogroupedMapFunction = Union[
 ArrowGroupedMapFunction = Union[
     Callable[[pyarrow.Table], pyarrow.Table],
     Callable[[Tuple[pyarrow.Scalar, ...], pyarrow.Table], pyarrow.Table],
+    Callable[[pyarrow.Table], Iterator[pyarrow.RecordBatch]],
+    Callable[[Tuple[pyarrow.Scalar, ...], pyarrow.Table], Iterator[pyarrow.RecordBatch]],

Review Comment:
   Shouldn't we have iterator -> iterator instead of table -> iterator?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "EnricoMi (via GitHub)" <gi...@apache.org>.
EnricoMi commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1774820310


##########
python/pyspark/worker.py:
##########
@@ -333,17 +336,38 @@ def wrap_cogrouped_map_arrow_udf(f, return_type, argspec, runner_conf):
             (col.name, to_arrow_type(col.dataType)) for col in return_type.fields
         ]
 
-    def wrapped(left_key_table, left_value_table, right_key_table, right_value_table):
-        if len(argspec.args) == 2:
-            result = f(left_value_table, right_value_table)
-        elif len(argspec.args) == 3:
-            key_table = left_key_table if left_key_table.num_rows > 0 else right_key_table
-            key = tuple(c[0] for c in key_table.columns)
-            result = f(key, left_value_table, right_value_table)
-
-        verify_arrow_result(result, _assign_cols_by_name, expected_cols_and_types)
+    def wrapped(left_key_batch, left_value_batches, right_key_batch, right_value_batches):
+        if is_generator:
+            if len(argspec.args) == 2:
+                result = f(left_value_batches, right_value_batches)
+            elif len(argspec.args) == 3:
+                key_batch = left_key_batch if left_key_batch.num_rows > 0 else right_key_batch
+                key = tuple(c[0] for c in key_batch.columns)
+                result = f(key, left_value_batches, right_value_batches)
+
+            def verify_element(batch):
+                verify_arrow_batch(batch, _assign_cols_by_name, expected_cols_and_types)
+                return batch
+
+            yield from map(verify_element, result)
+            # Make sure both iterators are fully consumed
+            for _ in left_value_batches:
+                pass
+            for _ in right_value_batches:
+                pass
 
-        return result.to_batches()
+        else:

Review Comment:
   We could get rid of this code duplication (the whole `False` branch) by turning `f` into an iterator function if `is_generator` is `False` (though the `len(argspec.args)` still duplicates some code).
   
   Add this right after the `import pyarrow as pa`:
   
       if not is_generator:
           def as_iterator(left: Iterator[pa.RecordBatch], right: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
               left_value_table = pa.Table.from_batches(left_value_batches)
               right_value_table = pa.Table.from_batches(right_value_batches)
               result = f(left_value_table, right_value_table)
               verify_arrow_table(result, _assign_cols_by_name, expected_cols_and_types)
               yield from result.to_batches()
   
           def as_iterator_with_key(key, left: Iterator[pa.RecordBatch], right: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]:
               left_value_table = pa.Table.from_batches(left_value_batches)
               right_value_table = pa.Table.from_batches(right_value_batches)
               result = f(key, left_value_table, right_value_table)
               verify_arrow_table(result, _assign_cols_by_name, expected_cols_and_types)
               yield from result.to_batches()
   
           if len(argspec.args) == 2:
               f = as_iterator
           elif len(argspec.args) == 3:
               f = as_iterator_with_key
   
   Then, we don't need to check `is_generated` in `wrapped`.



##########
python/pyspark/worker.py:
##########
@@ -333,17 +336,38 @@ def wrap_cogrouped_map_arrow_udf(f, return_type, argspec, runner_conf):
             (col.name, to_arrow_type(col.dataType)) for col in return_type.fields
         ]
 
-    def wrapped(left_key_table, left_value_table, right_key_table, right_value_table):
-        if len(argspec.args) == 2:
-            result = f(left_value_table, right_value_table)
-        elif len(argspec.args) == 3:
-            key_table = left_key_table if left_key_table.num_rows > 0 else right_key_table
-            key = tuple(c[0] for c in key_table.columns)
-            result = f(key, left_value_table, right_value_table)
-
-        verify_arrow_result(result, _assign_cols_by_name, expected_cols_and_types)
+    def wrapped(left_key_batch, left_value_batches, right_key_batch, right_value_batches):
+        if is_generator:
+            if len(argspec.args) == 2:
+                result = f(left_value_batches, right_value_batches)
+            elif len(argspec.args) == 3:
+                key_batch = left_key_batch if left_key_batch.num_rows > 0 else right_key_batch
+                key = tuple(c[0] for c in key_batch.columns)
+                result = f(key, left_value_batches, right_value_batches)
+
+            def verify_element(batch):
+                verify_arrow_batch(batch, _assign_cols_by_name, expected_cols_and_types)
+                return batch
+
+            yield from map(verify_element, result)
+            # Make sure both iterators are fully consumed

Review Comment:
   can we somehow close the iterators to stop transferring unused data?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2340315184

   Ok based on the feedback my plan is:
   - Make the grouped API iterator -> iterator (or (keys, iterator) -> iterator)
   - Make the cogrouped API (table, table) -> iterator (unless we think there's every going to be a way to split these into multiple batches. I would make them both iterators as well but will likely just always been an iterator of one element)
   - Keep the input iterator as an iterator of one element for now to get the API change in
   - As a follow on look into splitting the input into multiple batches based on arrowMaxRecordPerBatch
   
   > and I suggest we still keep existing Table -> Table interface, it should be able to support both.
   
   I would lean against this if possible as it introduces more possibility for user error and complicates the implementation a bit. Need to either check the annotation of the argument or choose based on whether the provided function is a generator or not. Not the end of the world, but given how easy it is for a user to convert from batches to a table and vice versa, seems like it might not be necessary?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2497891314

   Gentle ping again


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman closed pull request #48038: [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow
URL: https://github.com/apache/spark/pull/48038


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2505211865

   @Kimahriman Thanks for working on this, I personally prefer a PR containing both Python changes and Scala changes, to make sure the new feature works.
   I think this work could be split into several PRs in this way:
   
   - `grouped.applyInArrow` in Spark Classic;
   - `grouped.applyInArrow` in Spark Connect;
   - `grouped.applyInPandas` in Spark Classic;
   - `grouped.applyInPandas` in Spark Connect;
   - `cogrouped.applyInArrow` in Spark Classic;
   - `cogrouped.applyInArrow` in Spark Connect;
   - `cogrouped.applyInPandas` in Spark Classic;
   - `cogrouped.applyInPandas` in Spark Connect;
   
   And regarding `cogrouped.applyInArrow` with new iterator interface, I am still not sure how the RecordBatches of both sides should be split. Maybe we can focus on `grouped.applyInXXX` first.
   
   WDYT? @HyukjinKwon @xinrong-meng 


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "EnricoMi (via GitHub)" <gi...@apache.org>.
EnricoMi commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1774785181


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -803,27 +806,30 @@ def applyInArrow(
         Applies a function to each cogroup using Arrow and returns the result
         as a `DataFrame`.
 
-        The function should take two `pyarrow.Table`\\s and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and the two `pyarrow.Table`\\s.
-        For each side of the cogroup, all columns are passed together as a
-        `pyarrow.Table` to the user-function and the returned `pyarrow.Table` are combined as
-        a :class:`DataFrame`.
+        The function can take one of two forms: It can take two `pyarrow.Table`\\s and return a
+        `pyarrow.Table`, or it can take two iterators of `pyarrow.RecordBatch` and yield
+        `pyarrow.RecordBatch`. Alternatively, each form can take a tuple of `pyarrow.Scalar`
+        as the first argument in addition to the input type above. For each cogroup, all columns
+        are passed together in the `pyarrow.Table` or `pyarrow.RecordBatch`, and the returned
+        `pyarrow.Table` or iterator of `pyarrow.RecordBatch` are combined as a :class:`DataFrame`.
 
         The `schema` should be a :class:`StructType` describing the schema of the returned

Review Comment:
   We should add an implementation note here saying that
   
       The content of the first iterator of `pyarrow.RecordBatch` will be hold entirely in memory,
       while for the second iterator only the current `pyarrow.RecordBatch` is kept in memory.
       Users should use the dataframe that they anticipate to produce larger groups as the
       right argument of the cogroup to reduce memory footprint.
   
   That way, users can already benefit from the iterator interface. Loading only one of both groups into memory is an improvement over the table interface.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1774927552


##########
python/pyspark/worker.py:
##########
@@ -333,17 +336,38 @@ def wrap_cogrouped_map_arrow_udf(f, return_type, argspec, runner_conf):
             (col.name, to_arrow_type(col.dataType)) for col in return_type.fields
         ]
 
-    def wrapped(left_key_table, left_value_table, right_key_table, right_value_table):
-        if len(argspec.args) == 2:
-            result = f(left_value_table, right_value_table)
-        elif len(argspec.args) == 3:
-            key_table = left_key_table if left_key_table.num_rows > 0 else right_key_table
-            key = tuple(c[0] for c in key_table.columns)
-            result = f(key, left_value_table, right_value_table)
-
-        verify_arrow_result(result, _assign_cols_by_name, expected_cols_and_types)
+    def wrapped(left_key_batch, left_value_batches, right_key_batch, right_value_batches):
+        if is_generator:
+            if len(argspec.args) == 2:
+                result = f(left_value_batches, right_value_batches)
+            elif len(argspec.args) == 3:
+                key_batch = left_key_batch if left_key_batch.num_rows > 0 else right_key_batch
+                key = tuple(c[0] for c in key_batch.columns)
+                result = f(key, left_value_batches, right_value_batches)
+
+            def verify_element(batch):
+                verify_arrow_batch(batch, _assign_cols_by_name, expected_cols_and_types)
+                return batch
+
+            yield from map(verify_element, result)
+            # Make sure both iterators are fully consumed

Review Comment:
   It would be great if we could but I don't really think it's possible. The iterators just come from a stream of bytes over a channel, there's no way to tell the upstream that it's done with this single arrow stream. If someone comes up with a way in the future it could always be improved later.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2645720820

   Closing in favor of https://github.com/apache/spark/pull/49005 which includes the JVM side batching as well


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2367827752

   I have the whole serialization side working for sending batches from the JVM side. I can go ahead and merge that in here and update the description to include that behavior as well unless it's easier to review them separately


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2340481093

   @Kimahriman thank you so much for working on it!
   
   1, regarding the cogrouped API, I personally prefer 
   ```
   (Iterator[RecordBatch], Iterator[RecordBatch]) -> Iterator[RecordBatch]
   (keys, Iterator[RecordBatch], Iterator[RecordBatch]) -> Iterator[RecordBatch]
   ```
   
   because I guess taking `Table` as input won't fix this issue
   https://github.com/apache/spark/blob/8fbbcb09454cbd62fdc7514ad2365200e7b123fb/python/pyspark/sql/pandas/group_ops.py#L873-L877
   
   
   2, regarding the grouped API, I think `Table -> Table` is convenient for simple cases and it is also consistent with `applyInPandas` which supports `pd.DataFrame -> pd.DataFrame`
   
   


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2340653183

   >because I guess taking Table as input won't fix this issue
   
   I can make the API change for cogrouped as well, but I don't think it's really going to be possible to to fix the loading all data into memory issue. Both left and right have to be serialized over the same channel, so you can't partially load the left and the right at the same time, you need to fully load the left at least first AFAIK. It would just be cosmetic but always an iterator of a single RecordBatch already loaded into memory


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "Kimahriman (via GitHub)" <gi...@apache.org>.
Kimahriman commented on PR #48038:
URL: https://github.com/apache/spark/pull/48038#issuecomment-2338897734

   I could try to see what it takes to update the input. The alternative is we could update the Python API so that it takes an iterator of `RecordBatch`, but just always make it an iterator of one element with the current serialization process. That would provide the flexibility to update the serialization side into multiple batches later without being a breaking change.
   
   Still end up with the inconsistency with cogroups though, maybe that's just fine?


-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "xinrong-meng (via GitHub)" <gi...@apache.org>.
xinrong-meng commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1751421772


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -538,26 +538,28 @@ def applyInArrow(
         as a `DataFrame`.
 
         The function should take a `pyarrow.Table` and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and a `pyarrow.Table`.
-        For each group, all columns are passed together as a `pyarrow.Table`
-        to the user-function and the returned `pyarrow.Table` are combined as a
-        :class:`DataFrame`.
+        `pyarrow.Table` or an iterator of `pyarrow.RecordBatch. Alternatively,
+        the user can pass a function that takes a tuple of `pyarrow.Scalar` grouping
+        key(s) and a `pyarrow.Table`. For each group, all columns are passed together
+        as a`pyarrow.Table` to the user-function and the returned `pyarrow.Table` or iterator

Review Comment:
   nit: missing a space after "a"



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org


Re: [PR] [SPARK-49547][SQL][PYTHON] Support returning iterator of RecordBatches in applyInArrow [spark]

Posted by "xinrong-meng (via GitHub)" <gi...@apache.org>.
xinrong-meng commented on code in PR #48038:
URL: https://github.com/apache/spark/pull/48038#discussion_r1751420860


##########
python/pyspark/sql/pandas/group_ops.py:
##########
@@ -538,26 +538,28 @@ def applyInArrow(
         as a `DataFrame`.
 
         The function should take a `pyarrow.Table` and return another
-        `pyarrow.Table`. Alternatively, the user can pass a function that takes
-        a tuple of `pyarrow.Scalar` grouping key(s) and a `pyarrow.Table`.
-        For each group, all columns are passed together as a `pyarrow.Table`
-        to the user-function and the returned `pyarrow.Table` are combined as a
-        :class:`DataFrame`.
+        `pyarrow.Table` or an iterator of `pyarrow.RecordBatch. Alternatively,

Review Comment:
   nit: missing backtick of pyarrow.RecordBatch



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org


---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org