You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "Yicong-Huang (via GitHub)" <gi...@apache.org> on 2026/03/23 22:32:53 UTC

[PR] [SPARK-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace manual type coercion logic [spark]

Yicong-Huang opened a new pull request, #54967:
URL: https://github.com/apache/spark/pull/54967

   ### What changes were proposed in this pull request?
   
   Replace manual column-by-column type coercion with `ArrowBatchTransformer.enforce_schema` in three places:
   
   1. `ArrowStreamArrowUDTFSerializer.apply_type_coercion` in serializers.py
   2. `ArrowStreamArrowUDFSerializer.create_batch` in serializers.py
   3. `process_results` in worker.py (scalar Arrow iter UDF path)
   
   Also:
   - Add `arrow_cast` parameter to `enforce_schema` for strict type matching mode
   - Add `KeyError` handling in `enforce_schema` for missing columns with user-friendly error
   - Remove now-unused `coerce_arrow_array` imports from serializers.py and worker.py
   
   ### Why are the changes needed?
   
   These three places duplicated the same coerce-and-reassemble logic that `enforce_schema` already provides. Consolidating reduces code duplication and ensures consistent error handling.
   
   ### Does this PR introduce _any_ user-facing change?
   
   Error messages for type/schema mismatches in Arrow UDTFs are slightly changed to be consistent with other Arrow UDF error messages.
   
   ### How was this patch tested?
   
   Existing tests in `test_arrow_udtf.py` and `test_arrow_udf_scalar.py`.
   
   ### Was this patch authored or co-authored using generative AI tooling?
   
   Yes.


-- 
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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #54967:
URL: https://github.com/apache/spark/pull/54967#discussion_r2983478064


##########
python/pyspark/sql/conversion.py:
##########


Review Comment:
   Good question! I believe the answer is no, as `use_index` only controls top-level column lookup. When column names are in the same order it uses positional index for speed, otherwise it falls back to name-based lookup. 
   
   Nested type coercion is handled by pa.Array.cast(), which Arrow applies recursively, so nested struct/list fields are covered regardless of the use_index path. I added a unit test (`test_enforce_schema_nested_cast`) to verify 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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

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

   merged to master


-- 
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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

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


##########
python/pyspark/sql/conversion.py:
##########


Review Comment:
   on `use_index`, do we need to take care of nested cases?



-- 
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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

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


##########
python/pyspark/sql/conversion.py:
##########
@@ -149,8 +154,19 @@ def enforce_schema(
 
         coerced_arrays = []
         for i, field in enumerate(arrow_schema):
-            arr = batch.column(i) if use_index else batch.column(field.name)
+            try:
+                arr = batch.column(i) if use_index else batch.column(field.name)
+            except KeyError:

Review Comment:
   Nit: `batch.column(i)` (the `use_index=True` branch) would raise `IndexError`, not `KeyError`, if the index were ever out of bounds. Currently this is structurally unreachable because `use_index` requires identical name lists (hence identical lengths), but catching both would be more defensive and self-documenting about the two code paths.
   
   ```suggestion
               except (KeyError, IndexError):
   ```



-- 
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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

Posted by "Yicong-Huang (via GitHub)" <gi...@apache.org>.
Yicong-Huang commented on code in PR #54967:
URL: https://github.com/apache/spark/pull/54967#discussion_r2992276533


##########
python/pyspark/sql/conversion.py:
##########
@@ -149,8 +154,19 @@ def enforce_schema(
 
         coerced_arrays = []
         for i, field in enumerate(arrow_schema):
-            arr = batch.column(i) if use_index else batch.column(field.name)
+            try:
+                arr = batch.column(i) if use_index else batch.column(field.name)
+            except KeyError:

Review Comment:
   I think this would never happen? 



-- 
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-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic [spark]

Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng closed pull request #54967: [SPARK-56166][PYTHON] Use ArrowBatchTransformer.enforce_schema to replace column-wise type coercion logic
URL: https://github.com/apache/spark/pull/54967


-- 
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