diff --git a/backends-velox/src/main/scala/org/apache/gluten/execution/ColumnarPartialProjectExec.scala b/backends-velox/src/main/scala/org/apache/gluten/execution/ColumnarPartialProjectExec.scala index d657c3e3031..71936379286 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/execution/ColumnarPartialProjectExec.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/execution/ColumnarPartialProjectExec.scala @@ -297,6 +297,22 @@ object ColumnarPartialProjectExec { HiveUDFTransformer.isHiveUDF(h) && !VeloxHiveUDFTransformer.isSupportedHiveUDF(h) } + /** + * Returns true for generic catalyst expressions that Velox cannot offload and should be pulled + * into ColumnarPartialProject, excluding ScalaUDF/HiveUDF and non-row-evaluable wrappers. + */ + private def isUnmappedGlutenExpr(expr: Expression): Boolean = expr match { + case _: ScalaUDF => false + case _ if HiveUDFTransformer.isHiveUDF(expr) => false + case _: Literal | _: AttributeReference | _: BoundReference => false + case _: NamedExpression => false + case _ => + expr.deterministic && + !expr.isInstanceOf[Unevaluable] && + !expr.isInstanceOf[LambdaFunction] && + !ExpressionMappings.expressionsMap.contains(expr.getClass) + } + private def isBlacklistExpression(e: Expression): Boolean = { ExpressionMappings.blacklistExpressionMap.contains(e.getClass) } @@ -307,6 +323,7 @@ object ColumnarPartialProjectExec { case _: ScalaUDF => true case h if containsUnsupportedHiveUDF(h) => true case e if isBlacklistExpression(e) => true + case e if isUnmappedGlutenExpr(e) => true case p => p.children.exists(c => containsUDFOrBlacklistExpression(c)) } } @@ -339,6 +356,8 @@ object ColumnarPartialProjectExec { replaceByAlias(h, replacedAlias) case e if isBlacklistExpression(e) => replaceByAlias(e, replacedAlias) + case e if isUnmappedGlutenExpr(e) => + replaceByAlias(e, replacedAlias) case au @ Alias(_: ScalaUDF, _) => val replaceIndex = replacedAlias.indexWhere(r => r.exprId == au.exprId) if (replaceIndex == -1) { diff --git a/backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala index 41cd9ac186a..1ccf0b3fdf8 100644 --- a/backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala +++ b/backends-velox/src/test/scala/org/apache/gluten/expression/UDFPartialProjectSuite.scala @@ -20,9 +20,13 @@ import org.apache.gluten.config.GlutenConfig import org.apache.gluten.execution.{ColumnarPartialProjectExec, WholeStageTransformerSuite} import org.apache.spark.SparkConf +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, UnaryExpression} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.{ConstantFolding, NullPropagation} import org.apache.spark.sql.execution.ProjectExec import org.apache.spark.sql.functions.udf +import org.apache.spark.sql.types.{BooleanType, DataType} import java.io.File @@ -30,6 +34,17 @@ case class MyStruct(a: Long, b: Array[Long]) case class MyStructWithNullValue(a: Option[Long], b: Array[Long]) +case class DummyUnmappedExpr(child: Expression) extends UnaryExpression with CodegenFallback { + override def dataType: DataType = BooleanType + + override protected def nullSafeEval(input: Any): Any = { + input.asInstanceOf[Long] % 2 == 0 + } + + override protected def withNewChildInternal(newChild: Expression): DummyUnmappedExpr = + copy(child = newChild) +} + class UDFPartialProjectSuite extends WholeStageTransformerSuite { disableFallbackCheck override protected val resourcePath: String = "/tpch-data-parquet" @@ -64,9 +79,22 @@ class UDFPartialProjectSuite extends WholeStageTransformerSuite { spark.udf.register("no_argument", noArgument) val concat = udf((x: String) => x + "_concat") spark.udf.register("concat_concat", concat) + spark.sessionState.functionRegistry.registerFunction( + FunctionIdentifier("dummy_unmapped"), + new ExpressionInfo(classOf[DummyUnmappedExpr].getName, "dummy_unmapped"), + (children: Seq[Expression]) => DummyUnmappedExpr(children.head) + ) } + test("fallback generic unmapped expression via partial project") { + runQueryAndCompare( + "SELECT sum(cast(dummy_unmapped(cast(l_orderkey as long)) as int)" + + " + hash(l_partkey)) from lineitem") { + checkGlutenPlan[ColumnarPartialProjectExec] + } + } + ignore("test plus_one") { runQueryAndCompare("SELECT sum(plus_one(cast(l_orderkey as long))) from lineitem") { checkGlutenPlan[ColumnarPartialProjectExec]