[SPARK-55617][SQL] Add VariantGet to V2ExpressionBuilder for DSv2 filter pushdown

sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/VariantGet.java

@@ -0,0 +1,86 @@
+/*
+ * 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.connector.expressions;
+
+import java.util.Objects;
+
+import org.apache.spark.annotation.Evolving;
+import org.apache.spark.sql.internal.connector.ExpressionWithToString;
+import org.apache.spark.sql.types.DataType;
+
+/**
+ * Variant get expression.
+ *
+ * @since 4.1.0
+ */
+@Evolving
+public class VariantGet extends ExpressionWithToString {
+  private final Expression child;
+  private final String path;
+  private final DataType targetType;
+  private final boolean failOnError;
+  private final String timeZoneId;
+
+  /**
+   * Creates VariantGet expression.
+   * @param child variant column reference
+   * @param path JSON path string
+   * @param targetType expected result type
+   * @param failOnError whether to throw on cast failure ({@code variant_get}) or return null
+   *                    ({@code try_variant_get})
+   * @param timeZoneId timezone bound on the catalyst expression for timestamp casts, or null
+   */
+  public VariantGet(
+      Expression child,
+      String path,
+      DataType targetType,
+      boolean failOnError,
+      String timeZoneId) {
+    this.child = child;
+    this.path = path;
+    this.targetType = targetType;
+    this.failOnError = failOnError;
+    this.timeZoneId = timeZoneId;
+  }
+
+  public Expression child() { return child; }
+  public String path() { return path; }
+  public DataType targetType() { return targetType; }
+  public boolean failOnError() { return failOnError; }
+  public String timeZoneId() { return timeZoneId; }
+
+  @Override
+  public Expression[] children() { return new Expression[]{ child }; }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) return true;
+    if (o == null || getClass() != o.getClass()) return false;
+    VariantGet that = (VariantGet) o;
+    return failOnError == that.failOnError &&
+        Objects.equals(child, that.child) &&
+        Objects.equals(path, that.path) &&
+        Objects.equals(targetType, that.targetType) &&
+        Objects.equals(timeZoneId, that.timeZoneId);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(child, path, targetType, failOnError, timeZoneId);
+  }
+}

sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java

@@ -30,6 +30,7 @@
 import org.apache.spark.sql.connector.expressions.NamedReference;
 import org.apache.spark.sql.connector.expressions.GeneralScalarExpression;
 import org.apache.spark.sql.connector.expressions.GetArrayItem;
+import org.apache.spark.sql.connector.expressions.VariantGet;
 import org.apache.spark.sql.connector.expressions.Literal;
 import org.apache.spark.sql.connector.expressions.NullOrdering;
 import org.apache.spark.sql.connector.expressions.SortDirection;
@@ -113,6 +114,8 @@ public String build(Expression expr) {
         build(sortOrder.expression()), sortOrder.direction(), sortOrder.nullOrdering());
     } else if (expr instanceof GetArrayItem getArrayItem) {
       return visitGetArrayItem(getArrayItem);
+    } else if (expr instanceof VariantGet variantGet) {
+      return visitVariantGet(variantGet);
     } else if (expr instanceof GeneralScalarExpression e) {
       String name = e.name();
       if (isBinaryComparisonOperator(name)) {
@@ -419,6 +422,13 @@ protected String visitGetArrayItem(GetArrayItem getArrayItem) {
     );
   }
 
+  protected String visitVariantGet(VariantGet variantGet) {
+    throw new SparkUnsupportedOperationException(
+      "EXPRESSION_TRANSLATION_TO_V2_IS_NOT_SUPPORTED",
+      Map.of("expr", variantGet.toString())
+    );
+  }
+
   protected String visitExtract(Extract extract) {
     return visitExtract(extract.field(), build(extract.source()));
   }

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala

@@ -22,9 +22,10 @@ import org.apache.spark.internal.LogKeys.EXPR
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateFunction, Complete}
 import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
+import org.apache.spark.sql.catalyst.expressions.variant.VariantGet
 import org.apache.spark.sql.catalyst.optimizer.ConstantFolding
 import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
-import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, Extract => V2Extract, FieldReference, GeneralScalarExpression, GetArrayItem => V2GetArrayItem, LiteralValue, NullOrdering, SortDirection, SortValue, UserDefinedScalarFunc}
+import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, Extract => V2Extract, FieldReference, GeneralScalarExpression, GetArrayItem => V2GetArrayItem, LiteralValue, NullOrdering, SortDirection, SortValue, UserDefinedScalarFunc, VariantGet => V2VariantGet}
 import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Avg, Count, CountStar, GeneralAggregateFunc, Max, Min, Sum, UserDefinedAggregateFunc}
 import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And => V2And, Not => V2Not, Or => V2Or, Predicate => V2Predicate}
 import org.apache.spark.sql.internal.SQLConf
@@ -333,6 +334,18 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) extends L
         case _ =>
           None
       }
+    case v: VariantGet if v.path.foldable =>
+      (Option(v.path.eval()).map(_.toString), generateExpression(v.child)) match {
+        case (Some(path), Some(colRef: FieldReference)) =>
+          val vg = new V2VariantGet(colRef, path, v.targetType, v.failOnError,
+            v.timeZoneId.orNull)
+          if (isPredicate && v.dataType.isInstanceOf[BooleanType]) {
+            Some(new V2Predicate("BOOLEAN_EXPRESSION", Array[V2Expression](vg)))
+          } else {
+            Some(vg)
+          }
+        case _ => None
+      }
     // TODO supports other expressions
     case ApplyFunctionExpression(function, children) =>
       val childrenExpressions = children.flatMap(generateExpression(_))

sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ToStringSQLBuilder.scala

@@ -18,6 +18,7 @@
 package org.apache.spark.sql.internal.connector
 
 import org.apache.spark.sql.connector.expressions.GetArrayItem
+import org.apache.spark.sql.connector.expressions.VariantGet
 import org.apache.spark.sql.connector.util.V2ExpressionSQLBuilder
 
 /**
@@ -40,4 +41,13 @@ class ToStringSQLBuilder extends V2ExpressionSQLBuilder with Serializable {
   override protected def visitGetArrayItem(getArrayItem: GetArrayItem): String = {
     s"${getArrayItem.childArray.toString}[${getArrayItem.ordinal.toString}]"
   }
+
+  override protected def visitVariantGet(variantGet: VariantGet): String = {
+    val funcName = if (variantGet.failOnError()) "variant_get" else "try_variant_get"
+    val col = variantGet.child()
+    val path = variantGet.path()
+    val typ = variantGet.targetType().catalogString
+    val tz = Option(variantGet.timeZoneId()).map(z => s", tz=$z").getOrElse("")
+    s"$funcName($col, '$path', $typ$tz)"
+  }
 }

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala

@@ -21,12 +21,13 @@ import org.apache.spark.SparkConf
 import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.variant.VariantGet
 import org.apache.spark.sql.catalyst.util.V2ExpressionBuilder
-import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue}
+import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, GeneralScalarExpression, LiteralValue, VariantGet => V2VariantGet}
 import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And => V2And, Not => V2Not, Or => V2Or, Predicate}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
-import org.apache.spark.sql.types.{BooleanType, DoubleType, IntegerType, LongType, StringType, StructField, StructType}
+import org.apache.spark.sql.types.{BooleanType, DoubleType, IntegerType, LongType, StringType, StructField, StructType, TimestampType, VariantType}
 import org.apache.spark.unsafe.types.UTF8String
 
 class DataSourceV2StrategySuite extends SharedSparkSession {
@@ -818,6 +819,195 @@ class DataSourceV2StrategySuite extends SharedSparkSession {
         FieldReference("cdouble"))))
   }
 
+  test("VariantGet translates to V2VariantGet connector expression") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.city", StringType)
+    val expr = VariantGet(ref, path, StringType, failOnError = true)
+    val gt = GreaterThan(expr, Literal.create("NYC", StringType))
+    val result = new V2ExpressionBuilder(gt, isPredicate = true).build()
+    result match {
+      case Some(v2pred: Predicate) if v2pred.name() == ">" =>
+        v2pred.children()(0) match {
+          case vg: V2VariantGet =>
+            assert(vg.path() == "$.city")
+            assert(vg.targetType() == StringType)
+            assert(vg.failOnError())
+            assert(vg.timeZoneId() == null)
+            assert(vg.children().length == 1)
+            assert(vg.children()(0) == FieldReference("v"))
+          case other => fail(s"expected V2VariantGet, got ${other.getClass.getName}")
+        }
+      case _ => fail("expected predicate with name '>'")
+    }
+  }
+
+  test("try_variant_get translates with failOnError=false") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.city", StringType)
+    val expr = VariantGet(ref, path, StringType, failOnError = false)
+    val gt = GreaterThan(expr, Literal.create("NYC", StringType))
+    val result = new V2ExpressionBuilder(gt, isPredicate = true).build()
+    result match {
+      case Some(v2pred: Predicate) if v2pred.name() == ">" =>
+        v2pred.children()(0) match {
+          case vg: V2VariantGet =>
+            assert(!vg.failOnError())
+            assert(vg.path() == "$.city")
+          case other => fail(s"expected V2VariantGet, got ${other.getClass.getName}")
+        }
+      case _ => fail("expected predicate with name '>'")
+    }
+  }
+
+  test("VariantGet predicate is translated by translateFilterV2") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.city", StringType)
+    val expr = VariantGet(ref, path, StringType, failOnError = true)
+    val gt = GreaterThan(expr, Literal.create("NYC", StringType))
+    val result = DataSourceV2Strategy.translateFilterV2(gt)
+    assert(result.isDefined)
+    result.get.children()(0) match {
+      case vg: V2VariantGet =>
+        assert(vg.path() == "$.city")
+        assert(vg.targetType() == StringType)
+        assert(vg.failOnError())
+      case other =>
+        fail(s"expected V2VariantGet in translated predicate, got " +
+          s"${other.getClass.getName}")
+    }
+  }
+
+  test("VariantGet with integer targetType preserves type") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.count", StringType)
+    val expr = VariantGet(ref, path, IntegerType, failOnError = true)
+    val gt = GreaterThan(expr, Literal(100))
+    val result = new V2ExpressionBuilder(gt, isPredicate = true).build()
+    assert(result.isDefined)
+    result.get.children()(0) match {
+      case vg: V2VariantGet =>
+        assert(vg.path() == "$.count")
+        assert(vg.targetType() == IntegerType)
+      case other => fail(s"expected V2VariantGet, got ${other.getClass.getName}")
+    }
+  }
+
+  test("VariantGet with non-foldable path returns None") {
+    val ref = AttributeReference("v", VariantType)()
+    val s = AttributeReference("s", StringType)()
+    val expr = VariantGet(ref, s, StringType, failOnError = true)
+    val result = new V2ExpressionBuilder(expr).build()
+    assert(result.isEmpty, "non-foldable path should not translate")
+  }
+
+  test("VariantGet with foldable null path returns None") {
+    val ref = AttributeReference("v", VariantType)()
+    val nullPath = Literal.create(null, StringType)
+    val expr = VariantGet(ref, nullPath, StringType, failOnError = true)
+    val result = new V2ExpressionBuilder(expr).build()
+    assert(result.isEmpty, "null path should not translate (graceful, no NPE)")
+  }
+
+  test("VariantGet with non-column child returns None") {
+    val lit = Literal("v")
+    val path = Literal.create("$.a", StringType)
+    val expr = VariantGet(lit, path, StringType, failOnError = true)
+    val result = new V2ExpressionBuilder(expr).build()
+    assert(result.isEmpty, "non-column child should not translate")
+  }
+
+  test("VariantGet boolean targetType wraps in BOOLEAN_EXPRESSION predicate when isPredicate") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.flag", StringType)
+    val expr = VariantGet(ref, path, BooleanType, failOnError = true)
+    val result = new V2ExpressionBuilder(expr, isPredicate = true).build()
+    result match {
+      case Some(p: Predicate) if p.name() == "BOOLEAN_EXPRESSION" =>
+        p.children()(0) match {
+          case vg: V2VariantGet =>
+            assert(vg.targetType() == BooleanType)
+          case other =>
+            fail(s"expected V2VariantGet inside BOOLEAN_EXPRESSION, got " +
+              s"${other.getClass.getName}")
+        }
+      case _ => fail(s"expected BOOLEAN_EXPRESSION predicate, got $result")
+    }
+  }
+
+  test("VariantGet boolean targetType does not crash under Or (isPredicate path)") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.flag", StringType)
+    val boolExpr = VariantGet(ref, path, BooleanType, failOnError = true)
+    val x = AttributeReference("x", IntegerType)()
+    val orExpr = Or(boolExpr, GreaterThan(x, Literal(0)))
+    // A boolean-typed VariantGet in predicate position must translate to a V2Predicate, or the
+    // enclosing And/Or's `isInstanceOf[V2Predicate]` assert crashes planning;
+    // the BOOLEAN_EXPRESSION predicate provides that.
+    val result = new V2ExpressionBuilder(orExpr, isPredicate = true).build()
+    assert(result.isDefined, "Or with boolean VariantGet should translate without AssertionError")
+    result.get match {
+      case p: Predicate => // expected
+      case other => fail(s"expected a Predicate, got ${other.getClass.getName}")
+    }
+  }
+
+  test("VariantGet boolean targetType is scalar when not isPredicate") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.flag", StringType)
+    val expr = VariantGet(ref, path, BooleanType, failOnError = true)
+    val result = new V2ExpressionBuilder(expr, isPredicate = false).build()
+    result match {
+      case Some(vg: V2VariantGet) =>
+        assert(vg.targetType() == BooleanType)
+      case _ => fail(s"expected V2VariantGet scalar when isPredicate=false, got $result")
+    }
+  }
+
+  test("V2VariantGet toString renders as variant_get SQL") {
+    val ref = AttributeReference("v", VariantType)()
+    val vg = new V2VariantGet(FieldReference("v"), "$.city", StringType, true, null)
+    assert(vg.toString == "variant_get(v, '$.city', string)")
+  }
+
+  test("V2VariantGet toString renders as try_variant_get with timezone") {
+    val vg = new V2VariantGet(FieldReference("v"), "$.ts", TimestampType, false, "UTC")
+    assert(vg.toString == "try_variant_get(v, '$.ts', timestamp, tz=UTC)")
+  }
+
+  test("VariantGet with resolved timeZoneId passes it through the builder") {
+    val ref = AttributeReference("v", VariantType)()
+    val path = Literal.create("$.ts", StringType)
+    val expr = VariantGet(ref, path, TimestampType, failOnError = true, timeZoneId = Some("UTC"))
+    val gt = GreaterThan(expr, Literal.create(null, TimestampType))
+    val result = new V2ExpressionBuilder(gt, isPredicate = true).build()
+    assert(result.isDefined)
+    result.get.children()(0) match {
+      case vg: V2VariantGet =>
+        assert(vg.timeZoneId() == "UTC")
+        assert(vg.targetType() == TimestampType)
+      case other => fail(s"expected V2VariantGet, got ${other.getClass.getName}")
+    }
+  }
+
+  test("VariantGet with struct-nested variant column translates to nested FieldReference") {
+    val structType = StructType(Seq(StructField("v", VariantType)))
+    val parentRef = AttributeReference("s", structType)()
+    val nestedVariant = GetStructField(parentRef, 0)
+    val path = Literal.create("$.city", StringType)
+    val expr = VariantGet(nestedVariant, path, StringType, failOnError = true)
+    val gt = GreaterThan(expr, Literal.create("NYC", StringType))
+    val result = new V2ExpressionBuilder(gt, isPredicate = true).build()
+    assert(result.isDefined)
+    result.get.children()(0) match {
+      case vg: V2VariantGet =>
+        assert(vg.children()(0) == FieldReference(Seq("s", "v")))
+        assert(vg.path() == "$.city")
+        assert(vg.targetType() == StringType)
+      case other => fail(s"expected V2VariantGet with nested FieldReference, got " +
+        s"${other.getClass.getName}")
+    }
+  }
+
   test("Current Like functions are not supported") {
     val currentFunctions = Seq(
       CurrentDate(),