harshmotw-db commented on code in PR #57125:
URL: https://github.com/apache/spark/pull/57125#discussion_r3562652028
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala:
##########
@@ -1465,4 +1465,147 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
"VARIANT_SIZE_LIMIT",
Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`"))
}
+
+ test("variant_set") {
+ def checkSet(
+ input: String,
+ path: String,
+ value: Expression,
+ expected: String,
+ createIfMissing: Boolean = true): Unit = {
+ val expr = VariantSet(
+ Literal(parseJson(input)),
+ Literal.create(path, StringType),
+ value,
+ Literal(createIfMissing))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Object sets: replace an existing key, or add a missing one.
+ checkSet("""{"a": 1}""", "$.a", Literal(2), """{"a":2}""")
+ checkSet("""{"a": 1}""", "$.b", Literal(3), """{"a":1,"b":3}""")
+ checkSet("""{"a": {"b": 1}}""", "$.a.b", Literal(9), """{"a":{"b":9}}""")
+
+ // Missing intermediate keys are created when create_if_missing is true.
+ checkSet("{}", "$.a.b", Literal(1), """{"a":{"b":1}}""")
+ checkSet("{}", "$.a[0]", Literal(1), """{"a":[1]}""")
+
+ // Array sets replace the element at the index (no shifting).
+ checkSet("""["a","b","c"]""", "$[1]", Literal("z"), """["a","z","c"]""")
+ // N == length appends; N > length pads with variant nulls.
+ checkSet("""["a","b","c"]""", "$[3]", Literal("z"),
"""["a","b","c","z"]""")
+ checkSet("""["a","b"]""", "$[5]", Literal("z"),
"""["a","b",null,null,null,"z"]""")
+ checkSet("""{"a": [1, 2, 3]}""", "$.a[1]", Literal(9), """{"a":[1,9,3]}""")
+
+ // Array index as an intermediate segment.
+ checkSet("""[{"a": 1}]""", "$[0].b", Literal(2), """[{"a":1,"b":2}]""")
+ checkSet("""[{"a": 1}]""", "$[0].a", Literal(9), """[{"a":9}]""")
+ checkSet("""{"a": [1, 2]}""", "$.a[3].b", Literal(9),
"""{"a":[1,2,null,{"b":9}]}""")
+ // A missing intermediate array is created and padded with variant nulls
up to the index.
+ checkSet("{}", "$.a[2]", Literal(1), """{"a":[null,null,1]}""")
+
+ // scalastyle:off nonascii
+ // Non-ASCII (multi-byte UTF-8).
+ checkSet("""{"你好": 1}""", """$['你好']""", Literal(2), """{"你好":2}""")
+ checkSet("""{"a": 1}""", """$['世界']""", Literal(2), """{"a":1,"世界":2}""")
+ checkSet("{}", """$['café']['über']""", Literal("naïve"),
"""{"café":{"über":"naïve"}}""")
+ // scalastyle:on nonascii
+
+ // create_if_missing = false.
+ checkSet("""{"a": 1, "b": 2}""", "$.a", Literal(99), """{"a":99,"b":2}""",
false)
Review Comment:
Can you add more "good" cases for `create_if_default = false`? Right now, we
only have one good case, and I'd be interested in seeing replacements work in
nested inputs as well when `create_if_default = false`
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -942,6 +942,177 @@ object VariantInsert {
}
}
+case class VariantSet(
+ input: Expression,
+ path: Expression,
+ value: Expression,
+ createIfMissing: Expression)
+ extends QuaternaryExpression
+ with ExpectsInputTypes
+ with QueryErrorsBase {
+
+ override def first: Expression = input
+ override def second: Expression = path
+ override def third: Expression = value
+ override def fourth: Expression = createIfMissing
+
+ override def nullIntolerant: Boolean = true
+
+ override def dataType: DataType = VariantType
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ VariantType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ AnyDataType,
+ BooleanType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val result = super.checkInputDataTypes()
+ if (result.isFailure) {
+ result
+ } else if (value.dataType == NullType) {
+ TypeCheckResult.TypeCheckSuccess
+ } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps =
false)) {
+ DataTypeMismatch(
+ errorSubClass = "CAST_WITHOUT_SUGGESTION",
+ messageParameters =
+ Map("srcType" -> toSQLType(value.dataType), "targetType" ->
toSQLType(VariantType)))
+ } else {
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+
+ // When the path is a foldable expression, parse it once at planning time
and cache it. The Java
+ // segments are derived once per task (see `ParsedSetPath`), avoiding a
per-row conversion. `None`
+ // means the path is dynamic (or a foldable NULL, which makes the whole
expression NULL and is
+ // never evaluated).
+ @transient private lazy val foldablePath: Option[VariantSet.ParsedSetPath] =
{
+ if (path.foldable) {
+ val p = path.eval()
+ if (p == null) {
+ None
+ } else {
+ val s = p.asInstanceOf[UTF8String].toString
+ Some(VariantSet.ParsedSetPath(
+ VariantExpressionEvalUtils.parseVariantPath(s, prettyName), s))
+ }
+ } else {
+ None
+ }
+ }
+
+ override protected def nullSafeEval(v: Any, p: Any, valValue: Any, create:
Any): Any = {
+ val inputVariant = v.asInstanceOf[VariantVal]
+ val createIfMissingValue = create.asInstanceOf[Boolean]
+ foldablePath match {
+ case Some(parsed) =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, parsed.javaSegments, parsed.pathStr, valValue,
value.dataType,
+ createIfMissingValue, prettyName)
+ case None =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, p.asInstanceOf[UTF8String], valValue, value.dataType,
+ createIfMissingValue, prettyName)
+ }
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val cls = VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal, createVal) => {
+ val fromArg = ctx.addReferenceObj("from", value.dataType)
+ foldablePath match {
+ case Some(parsed) =>
+ val parsedArg = ctx.addReferenceObj("setPath", parsed)
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $parsedArg.javaSegments(), $parsedArg.pathStr(),
$valVal, $fromArg,
+ | $createVal, "$prettyName");
+ """.stripMargin
+ case None =>
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $pVal, $valVal, $fromArg, $createVal, "$prettyName");
+ """.stripMargin
+ }
+ })
+ }
+
+ override def prettyName: String = "variant_set"
+
+ override protected def withNewChildrenInternal(
+ newFirst: Expression,
+ newSecond: Expression,
+ newThird: Expression,
+ newFourth: Expression): VariantSet =
+ copy(input = newFirst, path = newSecond, value = newThird, createIfMissing
= newFourth)
+}
+
+object VariantSet {
+ // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
+ // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
+ // source string, retained for error messages.
+ case class ParsedSetPath(segments: Array[VariantPathSegment], pathStr:
String) {
+ @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
+ VariantExpressionEvalUtils.toJavaSegments(segments)
+ }
Review Comment:
I'm wondering if every variant construction "needs" to have its own
companion `Parsed***Path` class (for example, `ParsedInsertPath` and
`ParsedDeletePath`). Is there scope to reduce code duplication here?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala:
##########
@@ -1465,4 +1465,147 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
"VARIANT_SIZE_LIMIT",
Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`"))
}
+
+ test("variant_set") {
+ def checkSet(
+ input: String,
+ path: String,
+ value: Expression,
+ expected: String,
+ createIfMissing: Boolean = true): Unit = {
+ val expr = VariantSet(
+ Literal(parseJson(input)),
+ Literal.create(path, StringType),
+ value,
+ Literal(createIfMissing))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Object sets: replace an existing key, or add a missing one.
+ checkSet("""{"a": 1}""", "$.a", Literal(2), """{"a":2}""")
+ checkSet("""{"a": 1}""", "$.b", Literal(3), """{"a":1,"b":3}""")
+ checkSet("""{"a": {"b": 1}}""", "$.a.b", Literal(9), """{"a":{"b":9}}""")
+
+ // Missing intermediate keys are created when create_if_missing is true.
+ checkSet("{}", "$.a.b", Literal(1), """{"a":{"b":1}}""")
+ checkSet("{}", "$.a[0]", Literal(1), """{"a":[1]}""")
+
+ // Array sets replace the element at the index (no shifting).
+ checkSet("""["a","b","c"]""", "$[1]", Literal("z"), """["a","z","c"]""")
+ // N == length appends; N > length pads with variant nulls.
+ checkSet("""["a","b","c"]""", "$[3]", Literal("z"),
"""["a","b","c","z"]""")
+ checkSet("""["a","b"]""", "$[5]", Literal("z"),
"""["a","b",null,null,null,"z"]""")
+ checkSet("""{"a": [1, 2, 3]}""", "$.a[1]", Literal(9), """{"a":[1,9,3]}""")
+
+ // Array index as an intermediate segment.
+ checkSet("""[{"a": 1}]""", "$[0].b", Literal(2), """[{"a":1,"b":2}]""")
+ checkSet("""[{"a": 1}]""", "$[0].a", Literal(9), """[{"a":9}]""")
+ checkSet("""{"a": [1, 2]}""", "$.a[3].b", Literal(9),
"""{"a":[1,2,null,{"b":9}]}""")
+ // A missing intermediate array is created and padded with variant nulls
up to the index.
+ checkSet("{}", "$.a[2]", Literal(1), """{"a":[null,null,1]}""")
+
+ // scalastyle:off nonascii
+ // Non-ASCII (multi-byte UTF-8).
+ checkSet("""{"你好": 1}""", """$['你好']""", Literal(2), """{"你好":2}""")
+ checkSet("""{"a": 1}""", """$['世界']""", Literal(2), """{"a":1,"世界":2}""")
+ checkSet("{}", """$['café']['über']""", Literal("naïve"),
"""{"café":{"über":"naïve"}}""")
+ // scalastyle:on nonascii
+
+ // create_if_missing = false.
+ checkSet("""{"a": 1, "b": 2}""", "$.a", Literal(99), """{"a":99,"b":2}""",
false)
+ checkSet("""{"a": 1}""", "$.b", Literal(2), """{"a":1}""", false)
+ checkSet("""{"a": {"b": 1}}""", "$.a.c", Literal(2), """{"a":{"b":1}}""",
false)
+ checkSet("""["a","b"]""", "$[1]", Literal("z"), """["a","z"]""", false)
+ checkSet("""{"a": [1, 2]}""", "$.a[5]", Literal(9), """{"a":[1,2]}""",
false)
+
+ // Set a variant null vs. a verbatim string vs. structured JSON.
+ checkSet("""{"a": 1}""", "$.a", Literal(parseJson("null")),
"""{"a":null}""")
+ checkSet("{}", "$.a", Literal("""{"x":1}"""), """{"a":"{\"x\":1}"}""")
+ checkSet("{}", "$.a", Literal(parseJson("""{"x":1}""")),
"""{"a":{"x":1}}""")
+ checkSet(
+ "{}", "$.a", Literal.create(Array(1, 2, 3), ArrayType(IntegerType)),
"""{"a":[1,2,3]}""")
+
+ // NULL-intolerant: any NULL argument yields NULL.
+ checkEvaluation(
+ VariantSet(Literal.create(null, VariantType), Literal("$.a"),
Literal(1), Literal(true)),
+ null)
+ checkSet("""{"a": 1}""", null, Literal(1), null)
+ checkSet("""{"a": 1}""", "$.a", Literal.create(null, VariantType), null)
+ checkSet("""{"a": 1}""", "$.a", Literal.create(null, NullType), null)
+ checkEvaluation(
+ VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a"),
Literal(2),
+ Literal.create(null, BooleanType)),
+ null)
+
+ // Dynamic (non-foldable) path and create_if_missing.
+ val dynamic = VariantSet(
+ Literal(parseJson("""{"a": 1}""")),
+ BoundReference(0, StringType, nullable = true),
+ Literal(2),
+ BoundReference(1, BooleanType, nullable = true))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1,"b":2}""",
+ InternalRow(UTF8String.fromString("$.b"), true))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1}""",
+ InternalRow(UTF8String.fromString("$.b"), false))
+
+ // Type mismatch throws regardless of create_if_missing.
+ Seq(true, false).foreach { create =>
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$.a.b"),
Literal(2),
+ Literal(create)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_set`"))
+ }
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("5")), Literal("$.a"), Literal(2),
Literal(true)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
+
+ // Segment kind not matching the container: array index on an object,
object key on an array.
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("""{"a": 1}""")), Literal("$[0]"),
Literal(2), Literal(true)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$[0]", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("[1, 2]")), Literal("$.a"), Literal(2),
Literal(true)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_set`"))
+ // A key that needs bracket notation is rendered with brackets in
`failedAt`.
+ checkErrorInExpression[SparkRuntimeException](
+ VariantSet(Literal(parseJson("""{"a.b": 5}""")),
Literal("""$['a.b'].c"""), Literal(2),
+ Literal(true)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$['a.b'].c", "failedAt" -> "$['a.b']", "functionName" ->
"`variant_set`"))
+
+ // Structs and maps are rejected at analysis; arrays of scalars remain
allowed.
Review Comment:
Can you make sure converting them to variant using `to_variant_object` works
as expected? Not just `checkInputDataTypes = true` but also that the nested
fields are replaced by the variant values properly
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
