support GetStructField with scalar function as input

apache · Jan 24, 2025 · 20f345f · 20f345f
1 parent 291ff35
commit 20f345f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 3 deletions.
diff --git a/backends-velox/src/main/scala/org/apache/gluten/expression/ExpressionTransformer.scala b/backends-velox/src/main/scala/org/apache/gluten/expression/ExpressionTransformer.scala
@@ -56,7 +56,9 @@ case class VeloxGetStructFieldTransformer(
     child: ExpressionTransformer,
     ordinal: Int,
     original: GetStructField)
-  extends UnaryExpressionTransformer {
+  extends BinaryExpressionTransformer {
+  override def left: ExpressionTransformer = child
+  override def right: ExpressionTransformer = LiteralTransformer(ordinal)
   override def doTransform(args: Object): ExpressionNode = {
     val childNode = child.doTransform(args)
     childNode match {
@@ -70,8 +72,7 @@ case class VeloxGetStructFieldTransformer(
           node.getTypeNode.asInstanceOf[StructNode].getFieldTypes.get(ordinal)
         ExpressionBuilder.makeNullLiteral(nodeType)
       case _ =>
-        throw new GlutenNotSupportException(
-          s"Unsupported child expression of GetStructField: $original.")
+        super.doTransform(args)
     }
   }
 }

diff --git a/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/execution/MiscOperatorSuite.scala
@@ -1984,4 +1984,19 @@ class MiscOperatorSuite extends VeloxWholeStageTransformerSuite with AdaptiveSpa
       }
     }
   }
+
+  test("test GetStructField with scala function as input") {
+    withTable("t") {
+      Seq[String](
+        ("{\"a\":1,\"f\":20,\"e\":3}"),
+        ("{\"a\":1,\"f\":20}")
+      ).toDF().createOrReplaceTempView("t")
+      val query = "select from_json(a1, 'a INT, f INT, e INT').e from t"
+      runQueryAndCompare(query)(
+        df => {
+          val executedPlan = getExecutedPlan(df)
+          assert(executedPlan.count(_.isInstanceOf[ProjectExec]) == 0)
+        })
+    }
+  }
 }