dongjoon-hyun commented on PR #57899: URL: https://github.com/apache/spark/pull/57899#issuecomment-5241979278
Thank you for extending the SPARK-27052 lift to the vectorized scalar flavors. The overall design and wiring look solid and consistent with the existing conventions: the per-flavor eval types (103-106), `PythonUDF.SCALAR_TYPES`, the `canChainWithParallelUDFs` iterator constraint matching the worker's `num_udfs == 1` assert, `evalConf`'s `input_type` propagation via `isElementwiseUDF`, the typing stubs, and the docs are all in place, and the test coverage of null/empty arrays and multi-batch regrouping is nice. I found two bugs in the iterator elementwise path in `python/pyspark/worker.py`, plus the lint failure: **1. All arguments are flattened with the first argument's element type** The non-iter path correctly uses a per-offset type (`element_types[o]`), but the iter path computes a single `element_type = eval_conf.input_type[args_offsets[0]].dataType.elementType` and passes it to `_elementwise_flatten_arg` for every `o in args_offsets`. A multi-argument pandas iterator UDF whose argument types differ (e.g. `transform(int_arr, x -> iter_udf(x, str_col))` — reachable since the rewrite repeats an outer column into an aligned array) would have its second argument converted with the wrong Spark type, producing wrong results or a runtime error. The Arrow flavor is unaffected since `_elementwise_flatten_arg` ignores the type there. Could you use the per-offset element type like the non-iter path, and add a multi-argument iterator test? **2. `pa.concat_arrays` type mismatch for timestamp results in the pandas iterator flavor** `buffered` is initialized as `pa.array([], type=arrow_element_type)` where `arrow_element_type` comes from `to_arrow_type(..., timezone="UTC")`, but the pandas-flavor chunks come from `PandasToArrowConversion.convert(..., timezone=runner_conf.timezone)`, which types timestamps with the *session* timezone (`conversion.py`'s `convert_column`). For a `TIMESTAMP`-returning pandas iterator UDF (or a nested type containing one) with a non-UTC session timezone, the first `pa.concat_arrays([buffered, chunk])` raises `ArrowInvalid: arrays to be concatenated must be identically typed`. The Arrow flavor is consistent because `enforce_schema` coerces to the UTC-typed schema. Initializing the buffer with the session-timezone type for the pandas flavor (or lazily from the first chunk) would fix it; a timestamp-returning test would be good to have too. **3. Python linter CI failure** `ruff format` wants the `_elementwise_flatten_arg(batch.column(o), element_type, is_pandas, runner_conf)` call collapsed onto one line — same lines as fix (1), so they can be addressed together. Minor comments, take or leave: - In the non-iter path, `_elementwise_renest` (which builds the `ListArray`) runs before `verify_result_row_count`, so a wrong-length UDF result can surface as an opaque pyarrow offset error instead of the friendly length error. The 102 path verifies first; computing the total from the lengths before re-nesting would restore that. - The non-iter pandas flavor omits the base scalar-pandas path's "must return a `pandas.DataFrame` for `StructType`" check; the conversion still fails but with a less clear message. - The non-iter path re-flattens/converts argument columns per UDF; the 102 path flattens each column once per batch and shares it across fused UDFs. - `chunk.combine_chunks() if hasattr(...)` is dead code — `PandasToArrowConversion.convert` returns a `pa.RecordBatch`, so `.column(0)` is always a `pa.Array`; the "may be a ChunkedArray" comment is inaccurate. - The base iterator paths apply `verify_output_row_limit` as a fail-fast on over-production; the iter elementwise path only detects it at the end via `verify_result_row_count`, after the buffer has grown. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
