From 3a38ef1bb8c05d46816e18e8ca32b492d8de0fe8 Mon Sep 17 00:00:00 2001 From: Qiegang Long Date: Thu, 19 Feb 2026 22:30:50 -0500 Subject: [PATCH 1/3] [SPARK-55617] Add VariantGet to V2ExpressionBuilder for DSv2 filter pushdown Add a VariantGet case in V2ExpressionBuilder.generateExpression() so that variant_get and try_variant_get predicates can be translated into V2 UserDefinedScalarFunc and pushed down to connectors via SupportsPushDownV2Filters. This is to support file-level skipping for shredded variant columns in Iceberg. Only foldable paths and direct table column references for the variant column are supported. --- .../catalyst/util/V2ExpressionBuilder.scala | 16 ++++++ .../v2/DataSourceV2StrategySuite.scala | 56 ++++++++++++++++++- 2 files changed, 70 insertions(+), 2 deletions(-) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala index 72b466f5a0f9a..0d37a0b7ce7ab 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala @@ -22,6 +22,7 @@ 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} @@ -29,6 +30,7 @@ import org.apache.spark.sql.connector.expressions.aggregate.{AggregateFunc, Avg, 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 import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, StringType} +import org.apache.spark.unsafe.types.UTF8String /** * The builder to generate V2 expressions from catalyst expressions. @@ -333,6 +335,20 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) extends L case _ => None } + case v: VariantGet + if v.path.foldable + && v.child.isInstanceOf[Attribute] => + val colName = v.child.asInstanceOf[Attribute].name + val path = v.path.eval().toString + val typeName = v.dataType.catalogString + val colRef = FieldReference.column(colName) + val pathLit = LiteralValue(UTF8String.fromString(path), StringType) + val typeLit = LiteralValue(UTF8String.fromString(typeName), StringType) + val canonName = v.prettyName + Some(new UserDefinedScalarFunc( + canonName, + canonName, + Array[V2Expression](colRef, pathLit, typeLit))) // TODO supports other expressions case ApplyFunctionExpression(function, children) => val childrenExpressions = children.flatMap(generateExpression(_)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala index 5f89a618edd53..e7b8dce0dff98 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala +++ b/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, UserDefinedScalarFunc} 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, VariantType} import org.apache.spark.unsafe.types.UTF8String class DataSourceV2StrategySuite extends SharedSparkSession { @@ -818,6 +819,57 @@ class DataSourceV2StrategySuite extends SharedSparkSession { FieldReference("cdouble")))) } + test("VariantGet serializes to UserDefinedScalarFunc") { + 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 udf: UserDefinedScalarFunc => + assert(udf.name() == "variant_get") + assert(udf.children().length == 3) + case _ => fail("expected UserDefinedScalarFunc") + } + 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 udf: UserDefinedScalarFunc => + assert(udf.name() == "variant_get") + assert(udf.children().length == 3) + case _ => fail("expected UserDefinedScalarFunc in translated predicate") + } + } + + test("try_variant_get serializes to UserDefinedScalarFunc with try_variant_get name") { + 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 udf: UserDefinedScalarFunc => + assert(udf.name() == "try_variant_get") + assert(udf.children().length == 3) + case _ => fail("expected UserDefinedScalarFunc") + } + case _ => fail("expected predicate with name '>'") + } + } + test("Current Like functions are not supported") { val currentFunctions = Seq( CurrentDate(), From 7c169a3a67f7e260b69fed4131d48fd1e1012595 Mon Sep 17 00:00:00 2001 From: Qiegang Long Date: Thu, 11 Jun 2026 16:53:57 -0400 Subject: [PATCH 2/3] [SPARK-55617][SQL] Add VariantGet to V2ExpressionBuilder for DSv2 filter pushdown Address review comments. Major changes: - Use dedicated connector expression (VariantGet) instead of UDSF, which is cleaner and more correct. Follow GetArrayItem implementation. - Fix dropped timeZoneId and null-path NPE in VariantGet translation - Strengthen tests for VariantGet in DataSourceV2StrategySuite --- .../sql/connector/expressions/VariantGet.java | 86 ++++++++++ .../util/V2ExpressionSQLBuilder.java | 10 ++ .../catalyst/util/V2ExpressionBuilder.scala | 29 ++-- .../connector/ToStringSQLBuilder.scala | 10 ++ .../v2/DataSourceV2StrategySuite.scala | 147 +++++++++++++++--- 5 files changed, 244 insertions(+), 38 deletions(-) create mode 100644 sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/VariantGet.java diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/VariantGet.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/expressions/VariantGet.java new file mode 100644 index 0000000000000..35904fdf82288 --- /dev/null +++ b/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); + } +} diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java index 20ca3d2ac09e8..eaba2aa79df50 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/util/V2ExpressionSQLBuilder.java +++ b/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())); } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala index 0d37a0b7ce7ab..e40a6e6df88b5 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/util/V2ExpressionBuilder.scala @@ -25,12 +25,11 @@ 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 import org.apache.spark.sql.types.{BooleanType, DataType, IntegerType, StringType} -import org.apache.spark.unsafe.types.UTF8String /** * The builder to generate V2 expressions from catalyst expressions. @@ -335,20 +334,18 @@ class V2ExpressionBuilder(e: Expression, isPredicate: Boolean = false) extends L case _ => None } - case v: VariantGet - if v.path.foldable - && v.child.isInstanceOf[Attribute] => - val colName = v.child.asInstanceOf[Attribute].name - val path = v.path.eval().toString - val typeName = v.dataType.catalogString - val colRef = FieldReference.column(colName) - val pathLit = LiteralValue(UTF8String.fromString(path), StringType) - val typeLit = LiteralValue(UTF8String.fromString(typeName), StringType) - val canonName = v.prettyName - Some(new UserDefinedScalarFunc( - canonName, - canonName, - Array[V2Expression](colRef, pathLit, typeLit))) + 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(_)) diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ToStringSQLBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ToStringSQLBuilder.scala index 5c54f28976458..9df132418724d 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/connector/ToStringSQLBuilder.scala +++ b/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)" + } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala index e7b8dce0dff98..1edfeeaec6280 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala @@ -23,11 +23,11 @@ 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, UserDefinedScalarFunc} +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, VariantType} +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 { @@ -819,7 +819,7 @@ class DataSourceV2StrategySuite extends SharedSparkSession { FieldReference("cdouble")))) } - test("VariantGet serializes to UserDefinedScalarFunc") { + 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) @@ -828,10 +828,32 @@ class DataSourceV2StrategySuite extends SharedSparkSession { result match { case Some(v2pred: Predicate) if v2pred.name() == ">" => v2pred.children()(0) match { - case udf: UserDefinedScalarFunc => - assert(udf.name() == "variant_get") - assert(udf.children().length == 3) - case _ => fail("expected UserDefinedScalarFunc") + 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 '>'") } @@ -845,31 +867,112 @@ class DataSourceV2StrategySuite extends SharedSparkSession { val result = DataSourceV2Strategy.translateFilterV2(gt) assert(result.isDefined) result.get.children()(0) match { - case udf: UserDefinedScalarFunc => - assert(udf.name() == "variant_get") - assert(udf.children().length == 3) - case _ => fail("expected UserDefinedScalarFunc in translated predicate") + 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("try_variant_get serializes to UserDefinedScalarFunc with try_variant_get name") { + test("VariantGet with integer targetType preserves type") { 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 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(v2pred: Predicate) if v2pred.name() == ">" => - v2pred.children()(0) match { - case udf: UserDefinedScalarFunc => - assert(udf.name() == "try_variant_get") - assert(udf.children().length == 3) - case _ => fail("expected UserDefinedScalarFunc") + 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("expected predicate with name '>'") + 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))) + // Without the fix, And/Or assert V2Predicate and crash with AssertionError. + // With the fix, boolExpr is wrapped in BOOLEAN_EXPRESSION and Or translates. + 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("Current Like functions are not supported") { val currentFunctions = Seq( CurrentDate(), From 4929dcb2ef91df2f0d12f84b24e2d9f4899e4ffc Mon Sep 17 00:00:00 2001 From: Qiegang Long Date: Mon, 15 Jun 2026 12:49:56 -0400 Subject: [PATCH 3/3] [SPARK-55617][SQL] Add VariantGet to V2ExpressionBuilder for DSv2 filter pushdown Address review comments: - Fix "before/after" comment in test - Add two more tests for VariantGet --- .../v2/DataSourceV2StrategySuite.scala | 39 ++++++++++++++++++- 1 file changed, 37 insertions(+), 2 deletions(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala index 1edfeeaec6280..0301c1d0f5baa 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2StrategySuite.scala @@ -940,8 +940,9 @@ class DataSourceV2StrategySuite extends SharedSparkSession { val boolExpr = VariantGet(ref, path, BooleanType, failOnError = true) val x = AttributeReference("x", IntegerType)() val orExpr = Or(boolExpr, GreaterThan(x, Literal(0))) - // Without the fix, And/Or assert V2Predicate and crash with AssertionError. - // With the fix, boolExpr is wrapped in BOOLEAN_EXPRESSION and Or translates. + // 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 { @@ -973,6 +974,40 @@ class DataSourceV2StrategySuite extends SharedSparkSession { 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(),