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 72b466f5a0f9a..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 @@ -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(_)) 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 5f89a618edd53..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 @@ -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(),