dongjoon-hyun commented on code in PR #57899:
URL: https://github.com/apache/spark/pull/57899#discussion_r3754371680
##########
python/pyspark/worker.py:
##########
@@ -3122,6 +3222,257 @@ def func(split_index: int, data:
Iterator[pa.RecordBatch]) -> Iterator[pa.Record
# profiling is not supported for UDF
return func, None, ser, ser
+ if eval_type in (
+ PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ELEMENTWISE_UDF,
+ ):
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ require_minimum_pyarrow_version()
+
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ # A scalar pandas or Arrow UDF lifted out of a higher-order function's
lambda by
+ # ExtractPythonUDFFromLambda. Each argument arrives as ``array<T>``
aligned with the
+ # iterated array. We flatten each list column to its element column,
run the *vectorized*
+ # function once over that flat column (so it still receives a pandas
Series / DataFrame or a
+ # pa.Array, its native contract), then re-nest the flat result to
``array<R>`` using the
+ # input's offsets - one row in, one row out, one Python round trip per
batch.
+ is_pandas = eval_type ==
PythonEvalType.SQL_SCALAR_PANDAS_ELEMENTWISE_UDF
+
+ udf_infos = []
+ for udf_func, udf_args_offsets, udf_kwargs_offsets, udf_return_type in
udfs:
+ wrapped_func, args_kwargs_offsets = wrap_kwargs_support(
+ udf_func, udf_args_offsets, udf_kwargs_offsets
+ )
+ # The UDF returns one value per element, so its declared return
type is the element
+ # type of the ``array<R>`` this operator produces.
+ arrow_element_type = to_arrow_type(
+ udf_return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ udf_infos.append(
+ (wrapped_func, args_kwargs_offsets, udf_return_type,
arrow_element_type)
+ )
+ col_names = [f"_{i}" for i in range(len(udfs))]
+
+ if is_pandas:
+ import pandas as pd
+
+ # Each argument is ``array<T>``; the vectorized function must see the
element type ``T``.
+ element_types = [f.dataType.elementType for f in eval_conf.input_type]
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ for input_batch in data:
+ # Flatten each list column to its element array once per batch
and share it across
+ # fused UDFs (the 102 path does the same), rather than
re-flattening per UDF.
+ flat_columns = [
+ _elementwise_flatten_column(
+ col.flatten(), element_types[i], is_pandas, runner_conf
+ )
+ for i, col in enumerate(input_batch.itercolumns())
+ ]
+
+ output_arrays = []
+ for wrapped_func, offsets, return_type, arrow_element_type in
udf_infos:
+ # Re-nest by this UDF's first argument's shape. Different
UDFs in one operator
+ # may iterate differently shaped arrays, so each re-nests
by its own argument.
+ shape = input_batch.column(offsets[0])
+ shape_lengths = pc.list_value_length(shape).to_pylist()
+ total_elements = sum(n for n in shape_lengths if n is not
None)
+
+ result = wrapped_func(*[flat_columns[o] for o in offsets])
+ if is_pandas:
+ if not hasattr(result, "__len__"):
+ pd_type = (
+ "pandas.DataFrame"
+ if isinstance(return_type, StructType)
+ else "pandas.Series"
+ )
+ raise PySparkTypeError(
+ errorClass="UDF_RETURN_TYPE",
+ messageParameters={
+ "expected": pd_type,
+ "actual": type(result).__name__,
+ },
+ )
+ # struct return type must be a DataFrame (matches the
base pandas path).
+ if isinstance(return_type, StructType) and not
isinstance(
+ result, pd.DataFrame
+ ):
+ raise PySparkValueError(
+ "Invalid return type. Please make sure that
the UDF returns a "
+ "pandas.DataFrame when the specified return
type is StructType."
+ )
+ # Verify the flat length before re-nesting so a
wrong-length result raises the
+ # friendly RESULT_ROWS_MISMATCH rather than an opaque
pyarrow offset error.
+ verify_result_row_count(len(result), total_elements)
+
+ flat_arr = _elementwise_result_to_arrow(
+ result, return_type, arrow_element_type, is_pandas,
runner_conf
+ )
+ nested = _elementwise_renest(
+ flat_arr, shape_lengths,
pa.types.is_large_list(shape.type)
+ )
+ output_arrays.append(nested)
+
+ yield pa.RecordBatch.from_arrays(output_arrays, col_names)
+
+ # profiling is not supported for UDF
+ return func, None, ser, ser
+
+ if eval_type in (
+ PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF,
+ PythonEvalType.SQL_SCALAR_ARROW_ITER_ELEMENTWISE_UDF,
+ ):
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+
+ require_minimum_pyarrow_version()
+
+ import collections
+
+ import pyarrow as pa
+ import pyarrow.compute as pc
+
+ is_pandas = eval_type ==
PythonEvalType.SQL_SCALAR_PANDAS_ITER_ELEMENTWISE_UDF
+ if is_pandas:
+ import pandas as pd
+
+ assert num_udfs == 1, "One SCALAR_*_ITER_ELEMENTWISE UDF expected
here."
+ udf_func, args_offsets, kwargs_offsets, return_type = udfs[0]
+ assert not kwargs_offsets, "Iterator UDFs do not take keyword
arguments."
+
+ # A scalar iterator UDF (pandas or Arrow) lifted out of a higher-order
function's lambda.
+ # The user function keeps its iterator contract: it consumes an
iterator of batches and
+ # yields an iterator of batches, one output value per input value. We
preserve that by
+ # feeding it the *flattened* elements of each input batch and, since
the JVM joins UDF
+ # output to input positionally by row (one ``array<R>`` per input
``array<T>`` row, in
+ # order), buffering a FIFO of the per-row element counts to re-group
the streamed flat
+ # results back into arrays. Output batch boundaries need not match
input ones.
+ # Each argument is ``array<T>``; the pandas function must see each
argument's own element
+ # type ``T`` (arguments may differ, e.g. an outer column repeated into
an aligned array).
+ element_types = [f.dataType.elementType for f in eval_conf.input_type]
+ arrow_element_type = to_arrow_type(
+ return_type, timezone="UTC",
prefers_large_types=runner_conf.use_large_var_types
+ )
+ is_large = None # set from the first input batch; the list width is
uniform per column.
+
+ def func(split_index: int, data: Iterator[pa.RecordBatch]) ->
Iterator[pa.RecordBatch]:
+ # FIFO of per-row element counts (None for a null array) awaiting
their flat results,
+ # and the flat elements produced so far but not yet enough to
complete the head shapes.
+ pending_shapes: "collections.deque" = collections.deque()
+ num_input_elements = 0
+
+ def extract_flat(batch: pa.RecordBatch):
+ nonlocal is_large, num_input_elements
+ shape = batch.column(args_offsets[0])
+ if is_large is None:
+ is_large = pa.types.is_large_list(shape.type)
+ pending_shapes.append(pc.list_value_length(shape).to_pylist())
+ # Flatten each argument to its element column; the user
function sees the flat
+ # elements as a pandas Series / DataFrame (pandas) or a
pa.Array (Arrow). Each
+ # argument is converted with its own element type.
+ flat_cols = [
+ _elementwise_flatten_column(
+ batch.column(o).flatten(), element_types[o],
is_pandas, runner_conf
+ )
+ for o in args_offsets
+ ]
+ num_input_elements += len(flat_cols[0])
+ return flat_cols[0] if len(flat_cols) == 1 else
tuple(flat_cols)
+
+ flat_args_iter = map(extract_flat, data)
+
+ if not is_pandas:
+ verified_iter = verify_return_type(
+ udf_func(flat_args_iter),
+ Iterator[pa.Array], # type: ignore[type-abstract]
+ )
+ else:
+ pandas_iter_type = (
+ Iterator[pd.DataFrame]
+ if isinstance(return_type, StructType)
+ else Iterator[pd.Series]
+ )
+ verified_iter = verify_return_type(udf_func(flat_args_iter),
pandas_iter_type)
+
+ # Buffer the streamed flat element results and emit an
``array<R>`` row as soon as the
+ # shape at the head of the FIFO is fully covered. A row whose
length is 0 (an empty
+ # array) or None (a null array) needs no elements, so it is
emitted immediately even
+ # before any chunk arrives - this matters when a whole partition
is empty/null arrays
+ # and the UDF yields nothing, otherwise those rows would be
dropped by the positional
+ # JVM join. Chunks are held in a list and concatenated only when a
shape spans more than
+ # one, so a UDF that yields once per input batch (the common case)
never re-copies the
+ # buffer. ``empty_type`` supplies the element type for a
zero-length emit: the first
+ # chunk's type once seen (the pandas flavor types timestamps with
the session timezone),
+ # else the UTC-typed ``arrow_element_type``. A partition that only
ever emits
+ # zero-length rows never mixes the two, so the stream schema stays
consistent.
+ pending_chunks: "list" = []
+ pending_len = 0
+ empty_type = arrow_element_type
+ num_output_elements = 0
+
+ def emit_ready():
+ nonlocal pending_chunks, pending_len
+ while pending_shapes:
+ lengths = pending_shapes[0]
+ needed = sum(n for n in lengths if n is not None)
+ if needed > pending_len:
+ break
+ pending_shapes.popleft()
+ if needed == 0:
+ flat = pa.nulls(0, type=empty_type)
+ else:
+ combined = (
+ pending_chunks[0]
+ if len(pending_chunks) == 1
+ else pa.concat_arrays(pending_chunks)
+ )
+ flat = combined.slice(0, needed)
+ remainder = combined.slice(needed)
+ pending_chunks = [remainder] if len(remainder) else []
+ pending_len -= needed
+ nested = _elementwise_renest(flat, lengths, bool(is_large))
+ yield pa.RecordBatch.from_arrays([nested], ["_0"])
+
+ def process_results():
+ nonlocal pending_chunks, pending_len, empty_type,
num_output_elements
+ for result in verified_iter:
+ if is_pandas:
+ verify_pandas_result(
+ result,
+ return_type,
+ assign_cols_by_name=True,
+ truncate_return_schema=True,
+ )
+ chunk = _elementwise_result_to_arrow(
+ result, return_type, arrow_element_type, is_pandas,
runner_conf
+ )
+ num_output_elements += len(chunk)
+ # Fail fast if the UDF over-produces, before the buffer
grows unbounded (the
+ # base iterator paths do the same via
verify_output_row_limit).
+ if num_output_elements > num_input_elements:
+ raise PySparkRuntimeError(
+ errorClass="OUTPUT_EXCEEDS_INPUT_ROWS",
messageParameters={}
Review Comment:
`empty_type` is only updated for non-empty chunks, so a zero-length chunk
yielded for an input batch holding only empty/null arrays leaves it at the
UTC-typed `arrow_element_type`. The rows for that batch are emitted immediately
with that UTC type, fixing the output stream's schema; a later real chunk from
the pandas flavor is typed with the *session* timezone, so writing the next
output batch fails in `ArrowStreamSerializer.dump_stream` with
`pyarrow.lib.ArrowInvalid: Tried to write record batch with different schema`.
Reproduced with a timestamp-returning pandas iterator UDF, session timezone
`America/Los_Angeles`, `spark.sql.execution.arrow.maxRecordsPerBatch=1`, and a
single partition `[([],), ([1],)]` (see the suggested test in
`test_udf_in_higher_order_function.py`). A zero-length chunk still carries the
flavor's element type, so taking the type from every chunk keeps the stream
schema consistent. Verified the repro passes with this change and the existing
iterator tests (`test_scalar_iter_udf_over_all_empty_and_null_partition`,
`test_scalar_pandas_iter_udf_timestamp_return_type`, and the pandas/Arrow
iterator basics) still pass.
```suggestion
# Even a zero-length chunk carries the flavor's element
type (the pandas
# flavor types timestamps with the session timezone), so
always take it:
# otherwise rows emitted for an all-empty batch before
the first non-empty
# chunk would use the UTC-typed default and disagree
with later batches,
# breaking the output stream's single-schema contract.
empty_type = chunk.type
if len(chunk):
pending_chunks.append(chunk)
pending_len += len(chunk)
```
##########
python/pyspark/sql/tests/test_udf_in_higher_order_function.py:
##########
@@ -643,20 +643,218 @@ def test_udf_in_aggregate_still_fails(self):
df.select(agg).collect()
self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF",
str(ctx.exception))
- def test_pandas_udf_in_lambda_still_fails(self):
- # A vectorized (pandas) UDF in a lambda is not supported.
+ def test_scalar_pandas_udf_in_lambda(self):
+ # A vectorized scalar pandas UDF is lifted and applied over the
flattened elements, so it
+ # still receives a pandas Series (its native contract) once per batch,
not per element.
import pandas as pd
from pyspark.sql.functions import pandas_udf
- df = self.spark.createDataFrame([([1, 2],)], "values array<int>")
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>"
+ )
@pandas_udf(IntegerType())
def plus_one_pandas(s: pd.Series) -> pd.Series:
return s + 1
- with self.assertRaises(AnalysisException) as ctx:
- df.select(sf.transform("values", lambda x:
plus_one_pandas(x))).collect()
- self.assertIn("LAMBDA_FUNCTION_WITH_PYTHON_UDF", str(ctx.exception))
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_pandas(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, None],)],
+ )
+ # Composition around the UDF result, filter and array_sort key all
work too.
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: plus_one_pandas(x) *
2).alias("r")),
+ [([4, 6, 8],), ([],), (None,), ([22, None],)],
+ )
+
+ def test_scalar_arrow_udf_in_lambda(self):
+ # A vectorized scalar Arrow UDF is lifted the same way; it takes and
returns a pyarrow
+ # Array over the flattened elements.
+ import pyarrow as pa
+ from pyspark.sql.functions import arrow_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, None],)], "values array<int>"
+ )
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow(a: pa.Array) -> pa.Array:
+ return pa.compute.add(a, 1)
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_arrow(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, None],)],
+ )
+ assertDataFrameEqual(
+ df.select(sf.filter("values", lambda x: plus_one_arrow(x) >
2).alias("r")),
+ [([2, 3],), ([],), (None,), ([10],)],
+ )
+
+ def test_scalar_pandas_iter_udf_in_lambda(self):
+ # A scalar iterator pandas UDF keeps its iterator contract: it
consumes and produces an
+ # iterator of Series. The worker feeds it the flattened elements and
re-groups the streamed
+ # results back into arrays positionally, so output batch boundaries
need not match input.
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values
array<int>"
+ )
+
+ @pandas_udf(IntegerType())
+ def plus_one_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield s + 1
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_iter(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)],
+ )
+
+ def test_scalar_arrow_iter_udf_in_lambda(self):
+ # A scalar iterator Arrow UDF, lifted the same way as the pandas
iterator variant.
+ from typing import Iterator
+ import pyarrow as pa
+ from pyspark.sql.functions import arrow_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3],), ([],), (None,), ([10, 20],), ([5],)], "values
array<int>"
+ )
+
+ @arrow_udf(IntegerType())
+ def plus_one_arrow_iter(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
+ for a in it:
+ yield pa.compute.add(a, 1)
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x:
plus_one_arrow_iter(x)).alias("r")),
+ [([2, 3, 4],), ([],), (None,), ([11, 21],), ([6],)],
+ )
+
+ def test_scalar_pandas_iter_udf_multiple_arguments_differ_in_type(self):
+ # A two-argument iterator UDF whose arguments have different element
types: the array
+ # element (int) and an outer column (string) that the rewrite repeats
into an aligned
+ # array. Each argument must be flattened with its own element type,
not the first's.
+ from typing import Iterator, Tuple
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+
+ df = self.spark.createDataFrame(
+ [([1, 2, 3], "a"), ([], "b"), (None, "c"), ([10], "d")],
+ "values array<int>, tag string",
+ )
+
+ @pandas_udf(StringType())
+ def tag_each(it: Iterator[Tuple[pd.Series, pd.Series]]) ->
Iterator[pd.Series]:
+ for x, t in it:
+ yield t + x.astype("string")
+
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: tag_each(x,
sf.col("tag"))).alias("r")),
+ [(["a1", "a2", "a3"],), ([],), (None,), (["d10"],)],
+ )
+
+ def test_scalar_pandas_iter_udf_timestamp_return_type(self):
+ # A timestamp-returning pandas iterator UDF with a non-UTC session
timezone: the result
+ # chunks are typed with the session timezone, so the streamed buffer
must take its type
+ # from the first chunk rather than assuming UTC, or pa.concat_arrays
would fail. Assert both
+ # against the equivalent non-iterator pandas UDF (isolates the concat
fix) and against a
+ # native Spark expression computing the same instants (so a timezone
bug common to both UDF
+ # paths would still be caught, while going through identical
driver-collection semantics).
+ from typing import Iterator
+ import pandas as pd
+ from pyspark.sql.functions import pandas_udf
+ from pyspark.sql.types import TimestampType
+
+ with self.sql_conf({"spark.sql.session.timeZone":
"America/Los_Angeles"}):
+ df = self.spark.createDataFrame([([1, 2],), (None,), ([3],)],
"values array<int>")
+
+ def compute(x):
+ return pd.to_datetime(x, unit="D", origin="2020-01-01")
+
+ @pandas_udf(TimestampType())
+ def to_ts(s: pd.Series) -> pd.Series:
+ return compute(s)
+
+ @pandas_udf(TimestampType())
+ def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ yield compute(s)
+
+ iter_df = df.select(sf.transform("values", lambda x:
to_ts_iter(x)).alias("r"))
+ # Native equivalent: pandas interprets the tz-naive origin in the
session timezone, so
+ # `timestamp_add(DAY, x, TIMESTAMP '2020-01-01 00:00:00')` (also
session-local) matches.
+ native_df = df.select(
+ sf.transform(
+ "values",
+ lambda x: sf.timestamp_add("DAY", x,
sf.lit("2020-01-01").cast("timestamp")),
+ ).alias("r")
+ )
+ # Consistent with the non-iterator pandas UDF, and with the native
instants.
+ assertDataFrameEqual(
+ iter_df,
+ df.select(sf.transform("values", lambda x:
to_ts(x)).alias("r")),
+ )
+ assertDataFrameEqual(iter_df, native_df)
+
+ def test_scalar_iter_udf_over_all_empty_and_null_partition(self):
+ # SPARK-58695: when a whole partition holds only empty/null arrays,
the flattened inputs are
+ # all zero-length and a skip-empty iterator UDF yields no chunks.
Those rows still need one
+ # (empty / null) output row each, or the positional JVM join drops
them silently. Cover both
+ # the pandas and Arrow iterator flavors.
+ from typing import Iterator
+ import pandas as pd
+ import pyarrow as pa
+ from pyspark.sql.functions import pandas_udf, arrow_udf
+
+ # Single partition so the whole batch is empty/null arrays.
+ df = self.spark.createDataFrame(
+ [([],), (None,), ([],), (None,)], "values array<int>"
+ ).coalesce(1)
+
+ @pandas_udf(IntegerType())
+ def skip_empty_pandas(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
+ for s in it:
+ if len(s) == 0:
+ continue
+ yield s + 1
+
+ @arrow_udf(IntegerType())
+ def skip_empty_arrow(it: Iterator[pa.Array]) -> Iterator[pa.Array]:
+ for a in it:
+ if len(a) == 0:
+ continue
+ yield pa.compute.add(a, 1)
+
+ for f in (skip_empty_pandas, skip_empty_arrow):
+ assertDataFrameEqual(
+ df.select(sf.transform("values", lambda x: f(x)).alias("r")),
+ [([],), (None,), ([],), (None,)],
+ )
Review Comment:
Reproduction test for the `empty_type` issue flagged in `worker.py`: it
fails on the current revision with `ArrowInvalid: Tried to write record batch
with different schema` and passes with the suggested fix there.
```suggestion
)
def test_scalar_pandas_iter_udf_timestamp_after_empty_batch(self):
# A zero-length result chunk (from an input batch holding only
empty/null arrays) must not
# pin the output stream's timestamp type to the UTC-typed default:
the rows it emits and
# the rows emitted from a later real chunk (typed with the session
timezone) would then
# disagree, and the Arrow stream writer would reject the second
output batch.
import datetime
from typing import Iterator
import pandas as pd
from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import TimestampType
with self.sql_conf(
{
"spark.sql.session.timeZone": "America/Los_Angeles",
# One row per Arrow batch so the empty-array row forms its
own (first) batch.
"spark.sql.execution.arrow.maxRecordsPerBatch": "1",
}
):
df = self.spark.createDataFrame([([],), ([1],)], "values
array<int>").coalesce(1)
@pandas_udf(TimestampType())
def to_ts_iter(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
for s in it:
yield pd.to_datetime(s, unit="D", origin="2020-01-01")
assertDataFrameEqual(
df.select(sf.transform("values", lambda x:
to_ts_iter(x)).alias("r")),
[([],), ([datetime.datetime(2020, 1, 2)],)],
)
```
--
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]
