From 764adf614b0f3748bc9b35bc44585147a7549804 Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 10:49:34 +0100 Subject: [PATCH 1/2] perf: evaluate JVM-dispatched filter conjuncts only on rows passing preceding conjuncts DataFusion's AND evaluates its right side on every row once more than 20% of rows pass the left side. A conjunct that runs through the JVM codegen dispatcher (ScalaUDF, regex family outside the native subset, other dispatched expressions) then costs a JNI round trip and Arrow export for rows Spark would never evaluate it on. CometFilterExec now splits the serialized deterministic predicate into its top-level conjuncts and rewrites `a AND b`, where b contains a JvmScalarUdf, into `CASE WHEN a THEN b END`. DataFusion's CaseExpr projects the columns b needs, filters them by a and evaluates b on the selected rows only. Null results reject rows just like false, so the filter output is unchanged. Scan pushdown is untouched because only the native filter predicate changes. Controlled by spark.comet.exec.filter.shortCircuitJvmDispatch.enabled (default true). Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scala/org/apache/comet/CometConf.scala | 9 + .../apache/spark/sql/comet/operators.scala | 82 ++++++++- .../exec/CometFilterShortCircuitSuite.scala | 168 ++++++++++++++++++ 3 files changed, 256 insertions(+), 3 deletions(-) create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 03bd2701754..38f9129142a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -486,6 +486,15 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.filter.shortCircuitJvmDispatch.enabled") + .category(CATEGORY_EXEC) + .doc("Whether a native filter evaluates a top-level conjunct that runs through the JVM " + + "codegen dispatcher only on rows where all preceding conjuncts are true, by rewriting " + + "`a AND b` into `CASE WHEN a THEN b END`. Applies to deterministic conditions only.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index 0718697a794..1333c8c567f 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -53,13 +53,14 @@ import org.apache.spark.util.SerializableConfiguration import org.apache.spark.util.io.ChunkedByteBuffer import com.google.common.base.Objects -import com.google.protobuf.CodedOutputStream +import com.google.protobuf.{CodedOutputStream, Message} import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, ConfigEntry, ContribServices} import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.CometParquetUtils import org.apache.comet.rules.CometExecRule -import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.{CometOperatorSerde, Compatible, ExprOuterClass, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, isStringCollationType, supportedSortType} @@ -1428,15 +1429,90 @@ object CometFilterExec extends CometOperatorSerde[FilterExec] { val cond = exprToProto(op.condition, op.child.output) if (cond.isDefined && childOp.nonEmpty) { + val predicate = + if (CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.get() && + op.condition.deterministic) { + shortCircuitJvmDispatch(cond.get) + } else { + cond.get + } val filterBuilder = OperatorOuterClass.Filter .newBuilder() - .setPredicate(cond.get) + .setPredicate(predicate) Some(builder.setFilter(filterBuilder).build()) } else { None } } + def shortCircuitJvmDispatch(predicate: Expr): Expr = { + val conjuncts = flattenAnd(predicate) + val expensive = conjuncts.map(containsJvmScalarUdf) + if (!expensive.drop(1).contains(true)) { + predicate + } else { + val segments = ArrayBuffer(ArrayBuffer(conjuncts.head)) + conjuncts.zip(expensive).tail.foreach { case (conjunct, isExpensive) => + if (isExpensive) segments += ArrayBuffer(conjunct) else segments.last += conjunct + } + segments.tail.foldLeft(balancedAnd(segments.head.toIndexedSeq)) { (guard, segment) => + Expr + .newBuilder() + .setCaseWhen( + ExprOuterClass.CaseWhen + .newBuilder() + .addWhen(guard) + .addThen(balancedAnd(segment.toIndexedSeq))) + .build() + } + } + } + + private def flattenAnd(predicate: Expr): Seq[Expr] = { + val leaves = ArrayBuffer.empty[Expr] + var stack: List[Expr] = predicate :: Nil + while (stack.nonEmpty) { + val current = stack.head + stack = stack.tail + if (current.hasAnd) { + stack = current.getAnd.getLeft :: current.getAnd.getRight :: stack + } else { + leaves += current + } + } + leaves.toSeq + } + + private def balancedAnd(conjuncts: IndexedSeq[Expr]): Expr = { + if (conjuncts.length == 1) { + conjuncts.head + } else { + val mid = conjuncts.length / 2 + Expr + .newBuilder() + .setAnd( + ExprOuterClass.BinaryExpr + .newBuilder() + .setLeft(balancedAnd(conjuncts.slice(0, mid))) + .setRight(balancedAnd(conjuncts.slice(mid, conjuncts.length)))) + .build() + } + } + + private def containsJvmScalarUdf(message: Message): Boolean = message match { + case e: Expr if e.hasJvmScalarUdf => true + case _ => + message.getAllFields.values.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case children: java.util.List[_] => + children.asScala.exists { + case child: Message => containsJvmScalarUdf(child) + case _ => false + } + case _ => false + } + } + override def createExec(nativeOp: Operator, op: FilterExec): CometNativeExec = { CometFilterExec(nativeOp, op, op.output, op.condition, op.child, SerializedPlan(None)) } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala new file mode 100644 index 00000000000..e95e8455f72 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometFilterShortCircuitSuite.scala @@ -0,0 +1,168 @@ +/* + * 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.exec + +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.comet.{CometFilterExec, CometNativeScanExec, CometScanExec} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper + +import org.apache.comet.CometConf +import org.apache.comet.serde.ExprOuterClass.Expr + +class CometFilterShortCircuitSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val backreference = "^(\\\\d)\\\\1" + + override def beforeAll(): Unit = { + super.beforeAll() + spark.udf.register( + "strict_even", + (x: Int) => { + if (x % 3 != 0) throw new IllegalStateException(s"strict_even evaluated on $x") + x % 2 == 0 + }) + spark.udf.register("is_even", (x: Int) => x % 2 == 0) + } + + private def withData(f: => Unit): Unit = { + withTable("t") { + sql("""CREATE TABLE t USING parquet AS + |SELECT + | IF(id % 7 = 0, NULL, CAST(id AS INT)) AS x, + | IF(id % 5 = 0, NULL, id % 4 = 0) AS b, + | IF(id % 11 = 0, NULL, CAST(id * 37 % 1000 AS STRING)) AS s + |FROM range(4000)""".stripMargin) + f + } + } + + private def filterPredicates(df: DataFrame): Seq[Expr] = + collect(df.queryExecution.executedPlan) { case f: CometFilterExec => + f.nativeOp.getFilter.getPredicate + } + + private def caseWhenDepth(e: Expr): Int = + if (e.hasCaseWhen) 1 + caseWhenDepth(e.getCaseWhen.getWhen(0)) else 0 + + private def scanDataFilters(df: DataFrame): Seq[Expression] = + collect(df.queryExecution.executedPlan) { + case s: CometNativeScanExec => s.dataFilters + case s: CometScanExec => s.dataFilters + }.flatten + + private def checkRewritten(query: String, depth: Int = 1): Unit = { + val df = sql(query) + checkSparkAnswerAndOperator(df) + val predicates = filterPredicates(df) + assert(predicates.nonEmpty, df.queryExecution.executedPlan) + assert(predicates.map(caseWhenDepth).max == depth, predicates) + } + + private def checkNotRewritten(query: String): Unit = { + val df = sql(query) + checkSparkAnswer(df) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen), filterPredicates(df)) + } + + test("expensive conjunct is evaluated only on rows passing preceding conjuncts") { + withData { + checkRewritten("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + checkRewritten("SELECT x FROM t WHERE x > 10 AND x % 3 = 0 AND strict_even(x) AND x < 3000") + } + } + + test("disabled rewrite evaluates the expensive conjunct on every row") { + withData { + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql("SELECT x FROM t WHERE x % 3 = 0 AND strict_even(x)") + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + val e = intercept[Exception](df.collect()) + assert( + e.toString.contains("strict_even evaluated on") || + Option(e.getCause).exists(_.toString.contains("strict_even evaluated on"))) + } + } + } + + test("three-valued logic matches Spark") { + withData { + checkRewritten("SELECT x, b, s FROM t WHERE b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE NOT b AND is_even(x)") + checkRewritten("SELECT x, b, s FROM t WHERE (b OR x IS NULL) AND is_even(x)") + checkRewritten(s"SELECT x, b, s FROM t WHERE b AND s RLIKE '$backreference'") + checkRewritten(s"SELECT x, b, s FROM t WHERE x % 3 = 0 AND NOT (s RLIKE '$backreference')") + checkRewritten("SELECT x, b, s FROM t WHERE x % 2 = 0 AND (NOT is_even(x) OR b)") + } + } + + test("nested AND/OR and multiple expensive conjuncts") { + withData { + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND (b OR s RLIKE '$backreference') AND x > 100") + checkRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) AND (is_even(x) AND s RLIKE '$backreference')", + depth = 2) + checkRewritten( + s"SELECT x, s FROM t WHERE x % 3 = 0 AND strict_even(x) AND b AND s RLIKE '$backreference'", + depth = 2) + checkNotRewritten(s"SELECT x, s FROM t WHERE x % 3 = 0 OR s RLIKE '$backreference'") + checkNotRewritten( + s"SELECT x, s FROM t WHERE (x % 3 = 0 AND b) OR (is_even(x) AND s RLIKE '$backreference')") + } + } + + test("no rewrite without a preceding conjunct or JVM dispatch") { + withData { + checkNotRewritten("SELECT x FROM t WHERE is_even(x)") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND b AND s IS NOT NULL") + checkNotRewritten("SELECT x FROM t WHERE x % 3 = 0 AND s LIKE '1%'") + } + } + + test("nondeterministic condition is not rewritten") { + withData { + val df = sql( + "SELECT x FROM t WHERE monotonically_increasing_id() % 5 < 4 AND x % 3 = 0 AND is_even(x)") + assert(filterPredicates(df).nonEmpty, df.queryExecution.executedPlan) + assert( + filterPredicates(df).forall(p => !p.hasCaseWhen), + df.queryExecution.optimizedPlan.toString + df.queryExecution.executedPlan) + checkSparkAnswer(df) + } + } + + test("scan pushdown is unchanged") { + withData { + val query = "SELECT x FROM t WHERE x > 10 AND b AND is_even(x)" + val rewritten = sql(query) + var baselineFilters: Seq[Expression] = Nil + withSQLConf(CometConf.COMET_EXEC_FILTER_SHORT_CIRCUIT_JVM_DISPATCH_ENABLED.key -> "false") { + val df = sql(query) + assert(filterPredicates(df).forall(p => !p.hasCaseWhen)) + baselineFilters = scanDataFilters(df) + } + assert(filterPredicates(rewritten).exists(_.hasCaseWhen)) + assert(scanDataFilters(rewritten).nonEmpty) + assert(scanDataFilters(rewritten) == baselineFilters) + checkSparkAnswerAndOperator(rewritten) + } + } +} From 092d70af28b3ea4bbded9121335cfc54dbd07cfc Mon Sep 17 00:00:00 2001 From: msaf Date: Sat, 10 Oct 2026 12:52:00 +0100 Subject: [PATCH 2/2] ci: run CometFilterShortCircuitSuite Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + 2 files changed, 2 insertions(+) diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 1c11bd8de6d..473006aa378 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -529,6 +529,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 1cf2f06d213..d00c7f6d2e2 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -177,6 +177,7 @@ jobs: org.apache.comet.exec.CometFirstLastBoolAggFuzzSuite org.apache.comet.exec.CometExec3_4PlusSuite org.apache.comet.exec.CometExecSuite + org.apache.comet.exec.CometFilterShortCircuitSuite org.apache.comet.exec.CometTaskBinarySizeSuite org.apache.spark.sql.comet.CometNativeTaskKillSuite org.apache.comet.exec.CometEmptyRelationExecSuite