[SPARK-53323][CONNECT] Support df.asTable() for Arrow UDTF in Spark Connect #52320
+3
−5
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ueshin
approved these changes
Sep 12, 2025
Comment on lines
25
to
29
| def test_arrow_udtf_with_table_argument_basic(self): | ||
| super().test_arrow_udtf_with_table_argument_basic() | ||
|
|
||
| # TODO(SPARK-53323): Support table arguments in Spark Connect Arrow UDTFs | ||
| @unittest.skip("asTable() is not supported in Spark Connect") | ||
| def test_arrow_udtf_with_table_argument_and_scalar(self): | ||
| super().test_arrow_udtf_with_table_argument_and_scalar() |
There was a problem hiding this comment.
Let's remove these and add pass on the class body.
Comment on lines
+221
to
+224
| flattened_arrays = struct.flatten() | ||
| field_names = [field.name for field in struct.type] | ||
| flattened_batch = pa.RecordBatch.from_arrays( | ||
| struct.flatten(), schema=pa.schema(struct.type) | ||
| flattened_arrays, names=field_names |
There was a problem hiding this comment.
Hmm, but why it's working for Spark classic but not Spark Connect? This code path should be shared by both.
There was a problem hiding this comment.
I checked the generated schema from both with pyarrow 15.0.2 and 20.0.0:
print(f"{flattened_batch.schema} <=====> {pa.schema(struct.type)}")and saw from the tests:
id: int64 <=====> id: int64 not null
or
id: int64 <=====> id: int64
Seems like the original schema is more accurate?
HyukjinKwon
changed the title
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR supports df.asTable for Arrow UDTF in spark connect by correcting the schema creation in ArrowStreamArrowUDTFSerializer.load_stream(). The original code was incorrectly using pa.schema(struct.type) which created a schema with the entire struct as a single field, instead of extracting the individual field names with [field.name for field in struct.type] to properly create a RecordBatch from the flattened arrays.
This ensures table arguments are passed as proper pa.RecordBatch objects to Arrow UDTFs
Why are the changes needed?
This is a fix
Does this PR introduce any user-facing change?
No
How was this patch tested?
unit tests
Was this patch authored or co-authored using generative AI tooling?
No