andygrove commented on code in PR #4744:
URL: https://github.com/apache/datafusion-comet/pull/4744#discussion_r3760640865
##########
spark/src/main/scala/org/apache/comet/serde/arrays.scala:
##########
@@ -713,9 +713,32 @@ object CometFlatten extends CometExpressionSerde[Flatten]
with ArraysBase {
}
}
-object CometArrayFilter extends CometExpressionSerde[ArrayFilter] {
+object CometArrayFilter extends
CometHighOrderFunction[ArrayFilter]("array_filter") {
- override def getSupportLevel(expr: ArrayFilter): SupportLevel = Compatible()
+ private val UNARY_FUNCTION_EXPECTED =
+ "The array_filter function in DataFusion is limited to one lambda
parameter"
+
+ override def getUnsupportedReasons(): Seq[String] =
Seq(UNARY_FUNCTION_EXPECTED)
Review Comment:
This replaces the parent's reasons rather than adding to them, so `lambda
functions must be LambdaFunction` and `lambda arguments must be
NamedLambdaVariables` from `CometHighOrderFunction.getUnsupportedReasons()`
disappear from the generated compatibility docs for `filter`. Could this be
`super.getUnsupportedReasons() ++ Seq(UNARY_FUNCTION_EXPECTED)`?
##########
spark/src/main/scala/org/apache/comet/serde/CometHighOrderFunction.scala:
##########
@@ -0,0 +1,163 @@
+/*
+ * 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.
+ */
+
+package org.apache.comet.serde
+
+import scala.jdk.CollectionConverters._
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
HigherOrderFunction, LambdaFunction => SparkLambdaFunction, NamedLambdaVariable
=> SparkNamedLambdaVariable}
+
+import org.apache.comet.CometConf
+import org.apache.comet.CometSparkSessionExtensions.withFallbackReason
+import org.apache.comet.serde.CometHighOrderFunction.namedLambdaVariable2Proto
+import org.apache.comet.serde.ExprOuterClass.{HigherOrderFunc, LambdaFunction,
NamedLambdaVariable}
+import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal,
serializeDataType}
+
+/**
+ * Serializer that converts Spark higher-order functions (e.g. `filter`,
`transform`, `exists`)
+ * into Comet's protobuf representation.
+ *
+ * Depending on the available configuration and on whether the expression
satisfies the native
+ * constraints, [[convert]] produces one of two representations:
+ * - a native higher-order function proto (executed by the DataFusion
engine), used when
+ * `COMET_EXEC_HIGHER_ORDER_FUNCTION_NATIVE_ENABLED` is set and the
expression is natively
+ * supported (see [[nativeUnsupportedReason]] / [[getSupportLevel]]); or
+ * - a JVM codegen dispatch (Scala UDF fallback via
`CometScalaUDF.emitJvmCodegenDispatch`),
+ * used when the native path is unavailable but
`COMET_SCALA_UDF_CODEGEN_ENABLED` is enabled.
+ */
+case class CometHighOrderFunction[T <: HigherOrderFunction](name: String)
+ extends CometExpressionSerde[T] {
+
+ private val UNSUPPORTED_LAMBDA_TYPE = "lambda functions must be
LambdaFunction"
+ private val UNSUPPORTED_LAMBDA_PARAM_TYPE = "lambda arguments must be
NamedLambdaVariables"
+
+ override def getUnsupportedReasons(): Seq[String] =
+ Seq(UNSUPPORTED_LAMBDA_TYPE, UNSUPPORTED_LAMBDA_PARAM_TYPE)
+
+ private def nativeUnsupportedReason(expr: T): Option[String] = {
Review Comment:
I hit a case where the native path is taken but the lambda variable never
reaches the evaluated body, so the predicate is false for every element. This
is on fully default configuration, and it returns wrong results rather than
erroring.
```sql
CREATE TABLE t(sarr array<string>, a array<int>, b array<int>) USING parquet;
INSERT INTO t VALUES (array('abc','xyz','a1'), array(1,2,3),
array(10,20,30));
SELECT filter(sarr, x -> x rlike '^a') FROM t;
-- Spark: [abc, a1] Comet: []
SELECT filter(a, x -> exists(b, y -> y > x)) FROM t;
-- Spark: [1, 2, 3] Comet: []
SELECT filter(a, x -> array_max(transform(b, y -> y + x)) > 31) FROM t;
-- Spark: [2, 3] Comet: []
```
All three plan to `CometProject` and return empty arrays. Setting
`spark.comet.exec.higherOrderFunction.native.enabled=false` makes all three
correct, which points at the new native path.
My reading of the mechanism: `highOrderFunction2Proto` calls
`exprToProtoInternal` on the lambda body, and when the body contains `rlike`, a
non-native nested higher-order function like `exists` or `transform`, or
anything else that lands in `CometScalaUDF.emitJvmCodegenDispatch`, that call
succeeds and emits a `JvmScalarUdf` proto. So `hofProto.isDefined` is true and
we commit to native. But `emitJvmCodegenDispatch` binds against
`AttributeReference`s only, and a `NamedLambdaVariable` is not one, so the
per-element value never reaches the compiled kernel.
Could `nativeUnsupportedReason` walk the lambda body and decline the native
path when it finds a subexpression that will route through the codegen
dispatcher, including nested `HigherOrderFunction`s? Then `convert` degrades to
codegen the way the config doc describes.
Worth SQL tests for these shapes as well. One thing to watch: `filter(a, x
-> exists(b, y -> y > 15))` passes, because nothing crosses the lambda boundary
there. The capture has to be meaningful for the test to catch this.
##########
spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrayFilterBenchmark.scala:
##########
@@ -0,0 +1,87 @@
+/*
+ * 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.
+ */
+
+package org.apache.spark.sql.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+
+import org.apache.comet.CometConf
+
+// spotless:off
+/**
+ * Benchmark to measure performance of Comet array expressions. To run this
benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make
benchmark-org.apache.spark.sql.benchmark.CometArrayFilterBenchmark
+ * }}}
+ * Results will be written to
"spark/benchmarks/CometArrayFilterBenchmark-**results.txt".
+ */
+// spotless:on
+object CometArrayFilterBenchmark extends CometBenchmarkBase {
+
+ def runExprBenchmark(config: ArrayFilterExprConfig, values: Int, arraySize:
Int): Unit = {
+ val benchmark =
+ new Benchmark(s"${config.name} (size $arraySize)", values, output =
output)
+ withTempPath { dir =>
+ withTempTable("parquetV1Table") {
+ prepareTable(
+ dir,
+ spark.sql(
+ s"SELECT sequence(0, cast(rand(42) * $arraySize as int)) AS arr " +
+ s"FROM range($values)"))
+
+ benchmark.addCase(s"Spark ${config.name}") { _ =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark.sql(config.query).noop()
+ }
+ }
+
+ benchmark.addCase(s"Comet (Native) ${config.name}") { _ =>
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") {
+ spark.sql(config.query).noop()
+ }
+ }
+
+ benchmark.addCase(s"Comet (Codegen) ${config.name}") { _ =>
+ withSQLConf(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") {
+ spark.sql(config.query).noop()
+ }
+ }
+
+ benchmark.run()
+ }
+ }
+ }
+
+ def runCometBenchmark(args: Array[String]): Unit = {
+ val values = 4 * 1024 * 1024
+
+ val config =
+ ArrayFilterExprConfig("array_filter", "SELECT filter(arr, x -> x > 2)
FROM parquetV1Table")
+
+ runExprBenchmark(config, values, 100)
Review Comment:
The benchmark covers one shape only: `x -> x > 2` over int arrays of size
100. Since the stated rationale is avoiding the per-batch JNI call, it would be
more informative to also cover a predicate that captures an outer column and a
string-element case, so we can see whether the win holds across the shapes
people actually write.
--
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]
