cloud-fan commented on code in PR #56703:
URL: https://github.com/apache/spark/pull/56703#discussion_r3530562824
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala:
##########
@@ -1311,4 +1311,139 @@ class VariantExpressionSuite extends SparkFunSuite with
ExpressionEvalHelper {
noPaths.checkInputDataTypes()
}
}
+
+ test("variant_insert") {
+ def checkInsert(input: String, path: String, value: Expression, expected:
String): Unit = {
+ val expr = VariantInsert(
+ Literal(parseJson(input)), Literal.create(path, StringType), value)
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Object inserts.
+ checkInsert("""{"a": 1}""", "$.b", Literal(2), """{"a":1,"b":2}""")
+ checkInsert("{}", "$.a", Literal(1), """{"a":1}""")
+ checkInsert("""{"a": {"b": 1}}""", "$.a.c", Literal(3),
"""{"a":{"b":1,"c":3}}""")
+
+ // Missing intermediate keys are created; the next segment's kind picks
the container.
+ checkInsert("{}", "$.a.b.c", Literal(99), """{"a":{"b":{"c":99}}}""")
+ checkInsert("{}", "$.a[0]", Literal(1), """{"a":[1]}""")
+
+ // Array inserts shift existing elements right.
+ checkInsert("""["a","b","c"]""", "$[1]", Literal("z"),
"""["a","z","b","c"]""")
+ checkInsert("""["a","b","c"]""", "$[0]", Literal("z"),
"""["z","a","b","c"]""")
+ checkInsert("""["a","b","c"]""", "$[3]", Literal("z"),
"""["a","b","c","z"]""")
+ checkInsert("""["a","b"]""", "$[5]", Literal("z"),
"""["a","b",null,null,null,"z"]""")
+
+ // A missing intermediate created as an array is also padded with nulls up
to the index.
+ checkInsert("""{"a": []}""", "$.a[2].b", Literal(1),
"""{"a":[null,null,{"b":1}]}""")
+ checkInsert("""{"a": [1, 2]}""", "$.a[3].b", Literal(2),
"""{"a":[1,2,null,{"b":2}]}""")
+ checkInsert("{}", "$.a[2]", Literal(1), """{"a":[null,null,1]}""")
+ checkInsert("{}", "$.a[2].b", Literal(1), """{"a":[null,null,{"b":1}]}""")
+
+ // Descending through an existing container, then inserting into it.
+ checkInsert("""[{"a": 1}]""", "$[0].b", Literal(2), """[{"a":1,"b":2}]""")
+ checkInsert("""[{"x": 0}, {"a": 1}]""", "$[1].b", Literal(2),
"""[{"x":0},{"a":1,"b":2}]""")
+ checkInsert("""{"a": [1, 2, 3]}""", "$.a[1]", Literal(9),
"""{"a":[1,9,2,3]}""")
+
+ // Bracket and dot notations are interchangeable: `['k']`, `["k"]`, and
`.k` in one path.
+ checkInsert(
+ """{"a": {"b": {}}}""", """$['a'].b["c"]""", Literal(1),
"""{"a":{"b":{"c":1}}}""")
+
+ // Special / edge-case keys.
+ checkInsert("""{"a": 1}""", "$['']", Literal(2), """{"":2,"a":1}""")
+ checkInsert("""{"a": 1}""", "$['?']", Literal(2), """{"?":2,"a":1}""")
+ checkInsert(
+ """{"a": 1}""", "$['key with spaces']", Literal(2), """{"a":1,"key with
spaces":2}""")
+ checkInsert("""{"a": 1}""", "$.fb:testid", Literal(2),
"""{"a":1,"fb:testid":2}""")
+
+ // Insert a variant null (the JSON null literal) vs. a verbatim string vs.
structured JSON.
+ checkInsert("""{"a": 1}""", "$.b", Literal(parseJson("null")),
"""{"a":1,"b":null}""")
+ checkInsert("{}", "$.a", Literal("""{"x":1}"""), """{"a":"{\"x\":1}"}""")
+ checkInsert("{}", "$.a", Literal(parseJson("""{"x":1}""")),
"""{"a":{"x":1}}""")
+
+ // Values of various castable types exercise the corresponding
`castToVariant` branches.
+ checkInsert("{}", "$.a", Literal(true), """{"a":true}""")
+ checkInsert("{}", "$.a", Literal(7L), """{"a":7}""")
+ checkInsert("{}", "$.a", Literal(2.5), """{"a":2.5}""")
+ checkInsert(
+ "{}", "$.a", Literal.create(Array(1, 2, 3), ArrayType(IntegerType)),
"""{"a":[1,2,3]}""")
+
+ // NULL-intolerant: any NULL argument yields NULL.
+ checkEvaluation(
+ VariantInsert(Literal.create(null, VariantType), Literal("$.a"),
Literal(1)),
+ null)
+ checkInsert("""{"a": 1}""", null, Literal(1), null)
+ checkInsert("""{"a": 1}""", "$.b", Literal.create(null, VariantType), null)
+ checkInsert("""{"a": 1}""", "$.b", Literal.create(null, NullType), null)
+
+ val dynamic = VariantInsert(
+ Literal(parseJson("""{"a": 1}""")),
+ BoundReference(0, StringType, nullable = true),
+ Literal(2))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(dynamic, StringType)),
+ """{"a":1,"b":2}""",
+ InternalRow(UTF8String.fromString("$.b")))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$.a"),
Literal(2)),
+ "VARIANT_DUPLICATE_KEY",
+ Map("key" -> "a"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$.a.b"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 1}""")), Literal("$[0]"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$[0]", "failedAt" -> "$", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": [1, 2]}""")), Literal("$.a.b"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a.b", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a": 5}""")), Literal("$.a[0]"),
Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a[0]", "failedAt" -> "$.a", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("5")), Literal("$.a"), Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$.a", "failedAt" -> "$", "functionName" ->
"`variant_insert`"))
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("""{"a.b": 5}""")),
Literal("""$['a.b'].c"""), Literal(2)),
+ "VARIANT_PATH_TYPE_MISMATCH",
+ Map("path" -> "$['a.b'].c", "failedAt" -> "$['a.b']", "functionName" ->
"`variant_insert`"))
+
+ // Structs and maps are rejected at analysis.
+ Seq(
+ Literal.create(null, MapType(StringType, IntegerType)),
+ Literal.create(null, StructType(Seq(StructField("x", IntegerType))))
+ ).foreach { v =>
+ assert(
+ VariantInsert(Literal(parseJson("{}")), Literal("$.a"),
v).checkInputDataTypes().isFailure)
+ }
+
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("$"), Literal(1)),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "$", "functionName" -> "`variant_insert`"))
+ checkErrorInExpression[SparkRuntimeException](
+ VariantInsert(Literal(parseJson("{}")), Literal("abc"), Literal(1)),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "abc", "functionName" -> "`variant_insert`"))
+
+ val tooBig = "x".repeat(16 * 1024 * 1024)
Review Comment:
Optional (non-blocking): this `tooBig` case exercises `VARIANT_SIZE_LIMIT`
through a 16 MiB *value*, which trips `checkCapacity` inside `appendVariant`.
It doesn't cover the path the prior-round fix touched — a large intermediate
*array index* on a freshly-created container, where `appendNewPath`'s
null-padding loop is what must degrade to `VARIANT_SIZE_LIMIT`. A direct case
would lock that in:
```suggestion
val tooBig = "x".repeat(16 * 1024 * 1024)
checkErrorInExpression[SparkRuntimeException](
VariantInsert(Literal(parseJson("{}")), Literal("$.a[2000000000]"),
Literal(1)),
"VARIANT_SIZE_LIMIT",
Map("sizeLimit" -> "16.0 MiB", "functionName" -> "`variant_insert`"))
```
The existing `tooBig` assertion below can stay as-is; this just adds the
pad-index path alongside it.
--
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]
