andygrove opened a new pull request, #6714: URL: https://github.com/apache/datafusion-comet/pull/6714
## Which issue does this PR close? Closes #6706. ## Rationale for this change A typed Scala UDF with a boxed primitive parameter, such as `(x: java.lang.Long) => ...`, converts each value through Spark's input encoder, which runs a `SafeProjection` into a `GenericInternalRow` on every row. The issue measured this on #6697. The UDF it measured also returns `java.lang.Long`, and that result goes through the output encoder, which runs an `UnsafeProjection` into an `UnsafeRow` on every row and costs about as much as the input side. Timing the generated kernel alone on one 8192-row batch with a tenth of the rows null: | Kernel | ns per row | | --- | --- | | `(x: Long) => x + 1` | 2.8 | | `java.lang.Long` in and out, on `main` | 22.6 | | input encoder dropped | 15.2 | | output encoder dropped | 12.1 | | both dropped (this PR) | 7.2 | Dropping only the input encoder would close about a third of the gap, so this PR handles boxed results as well as boxed parameters. Spark's whole-stage codegen runs the same projections, so this is an improvement over Spark rather than a fix. ## What changes are included in this PR? `CometBatchKernelCodegen.generateSource` now rewrites each `ScalaUDF` in the tree before compiling it (`withoutBoxedPrimitiveEncoders`), setting the encoder of a boxed primitive parameter or result to `None`. With no encoder, `ScalaUDF` converts the value with `CatalystTypeConverters`, which passes a primitive Catalyst value through in its boxed form. Spark 4 builds a Java UDF the same way, with no encoders. The rewrite drops an encoder only when its expression shows that the conversion changes nothing: - For a parameter, the deserializer must be `<Boxed>.valueOf(input[0])` with null propagation, over a nullable reference of the matching primitive type, so a null input stays null. Over a non-nullable reference Spark would box the 0 that the row holds for a null, so that shape keeps its encoder. - For a result, the serializer must be `input[0].<primitive>Value()` on the boxed class. `ScalaUDF` turns a null result into null before it calls the serializer. Every other encoder (`String`, `Option`, case classes, collections) is left alone, so a UDF with mixed parameters loses only the encoders of its boxed ones. The `ScalaUDF`, `ExpressionEncoder`, `StaticInvoke` and `Invoke` code this relies on is the same in Spark 3.4.3, 3.5.8, 4.0.4, 4.1.3 and 4.2.0. End to end, this is `SELECT max(f(c))` over 4M Parquet rows with a tenth of them null. Each number is milliseconds above `max(c)` with no function, best of 15 runs. `main` and this PR ran interleaved in one JVM through a temporary switch, in a `CometTestBase` session with a release native build: | UDF | batch 1024 | batch 8192 | batch 65536 | | --- | --- | --- | --- | | `(x: Long) => x + 1` | 105 | 41 | 36 | | `java.lang.Long` in and out, `main` | 173 | 104 | 101 | | `java.lang.Long` in and out, this PR | 104 | 41 | 35 | | `java.lang.Long` in, `Long` out, `main` | 118 | 52 | 46 | | `java.lang.Long` in, `Long` out, this PR | 86 | 24 | 17 | A boxed UDF now costs the same as the primitive one. The boxed-in, primitive-out UDF costs less than either, because Spark puts no `if (isnull(c)) null else ...` guard around a boxed parameter, and that guard runs natively as a `CASE` (#6704). Setup: Apple M3 Max, JDK 17.0.10, Spark 4.1.3 (the default profile). #6709 adds `CometScalaUDFDispatchBenchmark` with these UDF shapes, so once both land it measures this change. ## How are these changes tested? Three new tests in `CometCodegenSuite`: - Which encoders the kernel drops. It drops both for each of the seven boxed types, only the parameter's for a boxed-in, primitive-out UDF, and none for primitive, `String` and `Option` parameters. A UDF with a boxed and a string parameter loses only the boxed one's. The test also checks the generated source: on the encoder path the kernel calls the parameter's converter even for a null input, and for a boxed parameter it no longer does. - The seven boxed types as parameters and results against Spark, over the minimum and maximum values, -0.0, NaN, the infinities and null. The `<type>_str` UDFs return the exact value they receive as a string (the raw bits for floats), so a null read as 0 or a -0.0 read as 0.0 would show. - Boxed parameters next to string, `Option` and primitive parameters, fed through Spark's cast of an int column, nested, and inside a `transform` lambda, where Spark evaluates the UDF through `eval` instead of generated code. `CometCodegenSuite`, `CometCodegenSourceSuite`, `CometCodegenHOFSuite` and `CometScalaUDFClassLoaderSuite` pass on Spark 3.4 (Scala 2.12), 4.1 and 4.2, and `CometCodegenFuzzSuite` passes on 4.1. Two mutations fail the new tests: dropping every input encoder fails the first test and the `Option` case, and skipping the rewrite in `generateSource` fails the source check. Spark's SQL suites have many UDF tests that would catch a behavior difference, so this PR has the `run-spark-4.1-tests` and `run-all-spark-profiles` labels. -- 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]
