andygrove opened a new pull request, #6365: URL: https://github.com/apache/datafusion-comet/pull/6365
## Which issue does this PR close? Closes #6290. ## Rationale for this change PySpark does not cast the batches a `mapInArrow` UDF returns to the declared output schema. It wraps them in a struct using their own Arrow schema, on 4.0, 4.1 and 4.2 alike. With `spark.comet.exec.pyarrowUDF.enabled=true`, `CometArrowPythonRunnerBase` wrapped each returned vector with `CometVector.getVector(vector, null)`, and `CometMapInBatchExec` read the columns at their declared types without checking them. `CometPlainVector` reads raw buffers at the declared width and scale. Running the shapes from the issue against the accelerated path on this branch's base (Spark 4.1, PySpark 4.1.3): - `int64` data under a declared `int`: `[0..7]` read as `[0, 0, 1, 0, 2, 0, 3, 0]`. - `int32` data under a declared `bigint`: read past the end of the buffer, returning values such as `4294967296` and `-2876713637903182319`. - `decimal(4,3)` data under a declared `decimal(10,2)`: `1.234` read as `12.34`. - Dictionary-encoded output under a declared `string`: `NullPointerException` from the null `DictionaryProvider`. Vanilla Spark 4.0, 4.1 and 4.2 raise for the integer and dictionary cases, because `ArrowColumnVector` picks its accessor from the physical vector. Spark rescales the decimal to `1.23`. With `spark.sql.execution.arrow.pyspark.validateSchema.enabled=true` (4.1+, off by default), `MapInBatchEvaluatorFactory` instead rejects any mismatch up front with `ARROW_TYPE_MISMATCH`. ## What changes are included in this PR? The runner now checks the output stream's schema once, when it reads the schema at `START_ARROW_STREAM`. That is before any batch is loaded or wrapped, and every batch in an Arrow stream shares that schema. On a mismatch it raises the same error Spark's `validateSchema` check raises: ``` [ARROW_TYPE_MISMATCH] Invalid schema from SQL_MAP_ARROW_ITER_UDF: expected StructType(StructField(id,IntegerType,true)), got StructType(StructField(id,LongType,true)). SQLSTATE: 42K0G ``` - `CometArrowPythonRunnerBase.outputSchemaMismatch` compares each returned Arrow field with the Arrow type Comet uses for the declared Spark type (`Utils.toArrowType`). It recurses through structs, lists and maps. - The comparison accepts only differences that do not change the values read, which Spark's `ArrowColumnVector` also reads the same way: field names, nullability, `large_string` and `large_binary` (which `CometPlainVector` reads), any time zone on a microsecond timestamp for `TimestampType`, collated strings, interval field ranges, a UDT's SQL type, and an all-null Arrow `null` column where the declared type allows nulls. Everything else is a mismatch, including dictionary-encoded columns, 256-bit decimals, `fixed_size_binary`, other timestamp units, and missing or extra columns. - The error describes the returned schema with Spark type names, as Spark does. It uses the Arrow schema when Spark type names cannot show the difference (for example `name: Int(32, true)[dictionary: 0]`). - The per-version runners build the error. Spark 4.1 and 4.2 use the `ARROW_TYPE_MISMATCH` condition. Spark 4.0 has no such condition, so it raises a `SparkException` with the same message. - `Spark4xMapInBatchSupport.RunnerInputs` carries the declared type (`PythonUDF.dataType`, as Spark's `MapInBatchExec` uses) to the runner through the per-version shims. - The PyArrow UDF user guide documents the check under Limitations. The same runner serves `mapInPandas`. There, PySpark converts the pandas output to the declared Arrow types, so the check passes and behavior is unchanged. No other Comet path reads Python output. Behavior differences from vanilla Spark with `validateSchema` off (the default): - Comet does not rescale decimals. It raises for a decimal whose precision or scale differs from the declared type, where Spark rescales. This matches Spark with `validateSchema` on. - Comet also raises for a timestamp without a time zone under `TimestampType`, for other types stored at the same width (for example `int64` under `timestamp`), and for extra columns. Spark's accessors read these. - With `validateSchema` on, Spark also rejects column name and nullability differences, which Comet accepts because they do not affect the values read. The dictionary `NullPointerException` goes away because dictionary-encoded output is now rejected before any vector is wrapped. Passing the reader as the `DictionaryProvider` would have made Comet return decoded strings where Spark raises, so this PR does not do that. ## How are these changes tested? - `CometArrowPythonRunnerSuite` (Spark 4.1 profile, 34 tests passed). New tests cover the integer width and decimal scale mismatches, including the expected messages, and layouts Comet would misread (dictionary, 256-bit decimal, `fixed_size_binary`, day-time interval, other timestamp units and zones, `large_list`, nested mismatches, `null` under a non-nullable type, missing or extra columns). They also cover differences that must still be accepted, and a reader-level test that feeds an IPC stream to the runner's output reader and gets the error at stream start. - `test_pyarrow_udf.py` adds `int64`-as-`int` and `int32`-as-`bigint` cases, the decimal scale case, the dictionary case, and a case with renamed columns, `large_string`/`large_binary`, a non-UTC time zone and an all-null `null` column that must read the same on both paths. Each runs in accelerated and fallback mode. The fallback runs confirm Spark raises for the integer and dictionary cases, rescales the decimal to `1.23`, and reads the last case. - Without the fix, the four accelerated mismatch tests failed (three did not raise, and the dictionary case raised `NullPointerException`). With the fix, locally with PySpark 4.1.3 against a `-Pspark-4.1` jar: `test_pyarrow_udf.py` 143 passed, and `test_pyarrow_udf_dictionary_shuffle.py` 6 passed. - Compiled main and test sources with `-Pspark-4.0` and `-Pspark-4.2`. PySpark 4.0 and 4.2 are left to the `run-pyarrow-udf-tests` CI job. -- 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]
