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 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) + } + } +}