diff --git a/.github/workflows/pyarrow_udf_test.yml b/.github/workflows/pyarrow_udf_test.yml index a3797f913b5..ce3737e0d8a 100644 --- a/.github/workflows/pyarrow_udf_test.yml +++ b/.github/workflows/pyarrow_udf_test.yml @@ -103,6 +103,8 @@ jobs: spark/src/test/resources/pyspark/test_pyarrow_udf.py /tmp/venv/bin/python -m pytest -v \ spark/src/test/resources/pyspark/test_pyarrow_udf_dictionary_shuffle.py + /tmp/venv/bin/python -m pytest -v \ + spark/src/test/resources/pyspark/test_pyarrow_udf_fuzz.py # Saved only on main: an entry written from a pull request or a # `gh-readonly-queue/*` branch cannot be restored by any later run, and diff --git a/dev/ci/compute-changes.py b/dev/ci/compute-changes.py index ae29077db4c..ef86766137d 100644 --- a/dev/ci/compute-changes.py +++ b/dev/ci/compute-changes.py @@ -157,6 +157,7 @@ "spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala", "spark/src/test/resources/pyspark/conftest.py", "spark/src/test/resources/pyspark/test_pyarrow_udf.py", + "spark/src/test/resources/pyspark/test_pyarrow_udf_fuzz.py", "spark/src/test/resources/pyspark/test_pyarrow_udf_dictionary_shuffle.py", "spark/src/test/spark-3.5/org/apache/spark/sql/comet/CometMapInBatchSuite.scala", "spark/src/test/spark-4.x/org/apache/spark/sql/comet/CometMapInBatchSuite.scala", diff --git a/spark/src/test/resources/pyspark/test_pyarrow_udf_fuzz.py b/spark/src/test/resources/pyspark/test_pyarrow_udf_fuzz.py new file mode 100644 index 00000000000..395a5c8fcd2 --- /dev/null +++ b/spark/src/test/resources/pyspark/test_pyarrow_udf_fuzz.py @@ -0,0 +1,379 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""Seeded data generation for PyArrow UDF fuzz coverage (issue #4384). + +The generator check runs without starting Spark. End-to-end tests consume the +same cases through Parquet and a real Python worker in both execution modes. +""" + +import datetime as dt +import inspect +import os +import random +from decimal import Decimal + +import pyarrow as pa +import pytest +from conftest import resolve_comet_jar +from pyspark.sql import Row, SparkSession, types as T +from pyspark.sql.pandas.types import to_arrow_schema + + +SEEDS = (4384, 4234, 5368) +NULL_FRACTIONS = (0.0, 0.01, 0.5, 0.99, 1.0) +MAX_DEPTH = 3 +COLLECTION_LENGTHS = (0, 1, 2, 7) +LARGE_ARRAY_LENGTH = 128 +VARIABLE_LENGTHS = (0, 1, 7, 31, 1024) + + +@pytest.fixture(scope="session") +def spark(): + jar = resolve_comet_jar() + os.environ["PYSPARK_SUBMIT_ARGS"] = ( + f"--jars {jar} --driver-class-path {jar} pyspark-shell" + ) + session = ( + SparkSession.builder.master("local[2]") + .appName("comet-pyarrow-udf-fuzz-tests") + .config("spark.plugins", "org.apache.spark.CometPlugin") + .config("spark.comet.enabled", "true") + .config("spark.comet.exec.enabled", "true") + # These tests scan one Parquet partition and do not require shuffle. + .config("spark.comet.shuffle.enabled", "false") + # Spark writes signed small integers. Allow ShortType scans despite the + # conservative guard for Parquet files containing unsigned UINT_8. + .config("spark.comet.scan.unsignedSmallIntSafetyCheck", "false") + .config("spark.sql.session.timeZone", "UTC") + .config("spark.memory.offHeap.enabled", "true") + .config("spark.memory.offHeap.size", "2g") + .getOrCreate() + ) + try: + yield session + finally: + session.stop() + + +def _primitive_types(rng): + precision = rng.choice((1, 9, 18, 19, 28, 38)) + return [ + T.BooleanType(), + T.ByteType(), + T.ShortType(), + T.IntegerType(), + T.LongType(), + T.FloatType(), + T.DoubleType(), + T.StringType(), + T.BinaryType(), + T.DateType(), + T.TimestampType(), + T.TimestampNTZType(), + T.DecimalType(precision, rng.randint(0, precision)), + ] + + +def _random_type(rng, depth, kind=None): + if depth == 0: + return rng.choice(_primitive_types(rng)) + kind = kind or rng.choice(("primitive", "array", "struct", "map")) + if kind == "primitive": + return rng.choice(_primitive_types(rng)) + if kind == "array": + return T.ArrayType(_random_type(rng, depth - 1), containsNull=True) + if kind == "struct": + return T.StructType( + [ + T.StructField(f"f{i}", _random_type(rng, depth - 1), nullable=True) + for i in range(rng.randint(1, 3)) + ] + ) + if kind == "map": + # Simple key types let us preserve map values as Python dictionaries. + return T.MapType( + rng.choice((T.StringType(), T.LongType())), + _random_type(rng, depth - 1), + valueContainsNull=True, + ) + raise ValueError(f"Unknown type category: {kind}") + + +def _random_value(rng, data_type, null_fraction): + if rng.random() < null_fraction: + return None + + if isinstance(data_type, T.BooleanType): + return bool(rng.getrandbits(1)) + for integer_type, bits in ( + (T.ByteType, 8), + (T.ShortType, 16), + (T.IntegerType, 32), + (T.LongType, 64), + ): + if isinstance(data_type, integer_type): + minimum, maximum = -(1 << (bits - 1)), (1 << (bits - 1)) - 1 + return rng.choice((minimum, maximum, 0, rng.randint(minimum, maximum))) + if isinstance(data_type, (T.FloatType, T.DoubleType)): + # Binary fractions in this range are exact even in float32. This keeps + # input rounding separate from the transport correctness being tested. + return rng.randint(-(1 << 20), 1 << 20) / 4.0 + if isinstance(data_type, T.DecimalType): + maximum = 10**data_type.precision - 1 + unscaled = rng.choice((0, maximum, rng.randint(0, maximum))) + digits = tuple(int(digit) for digit in str(unscaled)) + # Tuple construction is exact, including precision 38, regardless of + # Python's decimal context (whose default precision is only 28). + return Decimal((rng.getrandbits(1), digits, -data_type.scale)) + if isinstance(data_type, T.StringType): + return "".join( + rng.choice("aZ09 é中😀") for _ in range(rng.choice(VARIABLE_LENGTHS)) + ) + if isinstance(data_type, T.BinaryType): + return rng.randbytes(rng.choice(VARIABLE_LENGTHS)) + if isinstance(data_type, T.DateType): + return dt.date(2000, 1, 1) + dt.timedelta(days=rng.randint(-20000, 20000)) + if isinstance(data_type, (T.TimestampType, T.TimestampNTZType)): + return dt.datetime(2000, 1, 1) + dt.timedelta( + days=rng.randint(-20000, 20000), + seconds=rng.randrange(86400), + microseconds=rng.randrange(1000000), + ) + if isinstance(data_type, T.ArrayType): + lengths = COLLECTION_LENGTHS + # Keep outer containers small so large lengths do not multiply + # at every nesting level. + if not isinstance( + data_type.elementType, (T.ArrayType, T.StructType, T.MapType) + ): + lengths = (*lengths, LARGE_ARRAY_LENGTH) + return [ + _random_value(rng, data_type.elementType, null_fraction) + for _ in range(rng.choice(lengths)) + ] + if isinstance(data_type, T.StructType): + return { + field.name: _random_value(rng, field.dataType, null_fraction) + for field in data_type.fields + } + if isinstance(data_type, T.MapType): + # Sampling without replacement keeps keys unique and non-null. + keys = rng.sample(range(10000), rng.choice(COLLECTION_LENGTHS)) + if isinstance(data_type.keyType, T.StringType): + keys = [f"key_{key}" for key in keys] + return { + key: _random_value(rng, data_type.valueType, null_fraction) for key in keys + } + raise TypeError(f"Unsupported generated type: {data_type}") + +def _generate_case(seed, num_rows=37, null_fraction=None): + rng = random.Random(seed) + # Each case includes every primitive category plus each container category; + # recursive shapes, decimal types, column order and values depend on the seed. + data_types = _primitive_types(rng) + data_types.extend( + _random_type(rng, MAX_DEPTH, kind) for kind in ("array", "struct", "map") + ) + rng.shuffle(data_types) + fields = [T.StructField("row_id", T.LongType(), nullable=False)] + fields.extend( + T.StructField(f"c{i}", data_type) for i, data_type in enumerate(data_types) + ) + schema = T.StructType(fields) + null_fractions = [ + rng.choice(NULL_FRACTIONS) if null_fraction is None else null_fraction + for _ in data_types + ] + rows = [] + for row_id in range(num_rows): + row = {"row_id": row_id} + for field, fraction in zip(fields[1:], null_fractions): + row[field.name] = _random_value(rng, field.dataType, fraction) + rows.append(row) + return schema, rows + + +@pytest.mark.parametrize("seed", SEEDS, ids=lambda seed: f"seed-{seed}") +def test_generated_data_arrow_ipc_roundtrip(seed): + schema, rows = _generate_case(seed) + context = f"seed={seed}, schema={schema.json()}" + assert _generate_case(seed) == (schema, rows), context + + # Use UTC for TimestampType; TimestampNTZType remains timezone-naive. + # PySpark 4.2 replaced timestamp_utc with timezone. + if "timezone" in inspect.signature(to_arrow_schema).parameters: + arrow_schema = to_arrow_schema(schema, timezone="UTC") + else: + arrow_schema = to_arrow_schema(schema, timestamp_utc=True) + table = pa.Table.from_pylist(rows, schema=arrow_schema) + table.validate(full=True) + output = pa.BufferOutputStream() + with pa.ipc.new_stream(output, arrow_schema) as writer: + writer.write_table(table, max_chunksize=7) + with pa.ipc.open_stream(output.getvalue()) as reader: + restored = reader.read_all() + restored.validate(full=True) + assert restored.schema == arrow_schema, context + # Strict conversion also rejects any duplicate Arrow map keys; comparing + # ordered rows (with unique IDs) preserves multiplicity instead of using sets. + assert restored.to_pylist(maps_as_pydicts="strict") == table.to_pylist( + maps_as_pydicts="strict" + ), context + + +def _normalize(value): + if isinstance(value, Row): + value = value.asDict() + if isinstance(value, dict): + return {key: _normalize(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_normalize(item) for item in value] + if isinstance(value, bytearray): + return bytes(value) + return value + + +def _assert_rows_equal(actual, expected, context): + actual = sorted((_normalize(row) for row in actual), key=lambda row: row["row_id"]) + expected = sorted((_normalize(row) for row in expected), key=lambda row: row["row_id"]) + # Check the complete ID sequence, including multiplicity, before comparing + # fields. Sorting only IDs avoids imposing an ordering on nested maps. + assert [row["row_id"] for row in actual] == [ + row["row_id"] for row in expected + ], context + for actual_row, expected_row in zip(actual, expected): + assert actual_row.keys() == expected_row.keys(), context + for name, value in expected_row.items(): + assert actual_row[name] == value, ( + f"{context}, row_id={expected_row['row_id']}, column={name}" + ) + + +@pytest.mark.parametrize("seed", SEEDS, ids=lambda seed: f"seed-{seed}") +@pytest.mark.parametrize("batch_size", (7, 16), ids=lambda size: f"batch-{size}") +@pytest.mark.parametrize( + "null_fraction", + (None, *NULL_FRACTIONS), + ids=lambda fraction: "null-mixed" if fraction is None else f"null-{fraction}", +) +def test_map_in_arrow_randomized_roundtrip( + spark, tmp_path, seed, batch_size, null_fraction +): + schema, rows = _generate_case(seed, null_fraction=null_fraction) + context = ( + f"seed={seed}, batch_size={batch_size}, " + f"null_fraction={null_fraction}, schema={schema.json()}" + ) + source_path = str(tmp_path / "input.parquet") + settings = ( + "spark.comet.enabled", + "spark.comet.exec.pyarrowUDF.enabled", + "spark.comet.batchSize", + "spark.sql.execution.arrow.maxRecordsPerBatch", + ) + # Comet has defaults for these keys, but Spark RuntimeConfig.get does not + # expose a Comet default when the key has never been explicitly set. + previous = {key: spark.conf.get(key, None) for key in settings} + try: + # Use Spark to establish the Parquet baseline independently of Comet. + spark.conf.set("spark.comet.enabled", "false") + spark.createDataFrame(rows, schema).coalesce(1).write.parquet(source_path) + baseline = spark.read.parquet(source_path).collect() + _assert_rows_equal(baseline, rows, f"{context}, mode=parquet-baseline") + + spark.conf.set("spark.comet.enabled", "true") + # The accelerated path sends Comet source batches directly to Python; + # Spark's Arrow limit alone would not force multiple input batches. + spark.conf.set("spark.comet.batchSize", str(batch_size)) + spark.conf.set("spark.sql.execution.arrow.maxRecordsPerBatch", str(batch_size)) + + def annotate_batches(iterator): + for batch_index, batch in enumerate(iterator): + batch.validate(full=True) + yield pa.RecordBatch.from_arrays( + [ + *batch.columns, + pa.array([batch_index] * batch.num_rows, type=pa.int32()), + pa.array([batch.num_rows] * batch.num_rows, type=pa.int32()), + ], + names=[*batch.schema.names, "batch_index", "batch_size"], + ) + + expected_batch_sizes = [ + min(batch_size, len(rows) - start) + for start in range(0, len(rows), batch_size) + ] + + outputs = {} + for accelerated in (True, False): + mode = "accelerated" if accelerated else "fallback" + spark.conf.set( + "spark.comet.exec.pyarrowUDF.enabled", str(accelerated).lower() + ) + source = spark.read.parquet(source_path) + # Batch indices restart per partition; this bounded input is written + # into one small file so all batches come from one worker iterator. + assert source.rdd.getNumPartitions() == 1, f"{context}, mode={mode}" + output_schema = T.StructType( + [ + *source.schema.fields, + T.StructField("batch_index", T.IntegerType(), nullable=False), + T.StructField("batch_size", T.IntegerType(), nullable=False), + ] + ) + result = source.mapInArrow(annotate_batches, output_schema) + plan = result._jdf.queryExecution().executedPlan().toString() + if accelerated: + assert "CometMapInBatch" in plan, f"{context}, mode={mode}\n{plan}" + assert "ColumnarToRow" not in plan, f"{context}, mode={mode}\n{plan}" + else: + assert "CometMapInBatch" not in plan, f"{context}, mode={mode}\n{plan}" + assert "MapInArrow" in plan, f"{context}, mode={mode}\n{plan}" + annotated_rows = result.collect() + observed_batches = {} + for row in annotated_rows: + observed_batches.setdefault(row["batch_index"], []).append(row) + assert sorted(observed_batches) == list(range(len(expected_batch_sizes))), ( + f"{context}, mode={mode}, observed_batches={sorted(observed_batches)}" + ) + for index, expected_size in enumerate(expected_batch_sizes): + batch_rows = observed_batches[index] + batch_context = f"{context}, mode={mode}, batch_index={index}" + assert len(batch_rows) == expected_size, batch_context + assert {row["batch_size"] for row in batch_rows} == { + expected_size + }, batch_context + + # Compare every original field after checking worker batch metadata. + # Full row IDs preserve multiplicity and catch losses/duplicates at + # batch boundaries as well as changes to nested payloads. + outputs[mode] = [ + {name: row[name] for name in source.columns} for row in annotated_rows + ] + _assert_rows_equal(outputs[mode], baseline, f"{context}, mode={mode}") + + _assert_rows_equal( + outputs["accelerated"], outputs["fallback"], f"{context}, mode=cross-check" + ) + finally: + for key, value in previous.items(): + if value is None: + spark.conf.unset(key) + else: + spark.conf.set(key, value)