chenhao-db commented on code in PR #57827:
URL: https://github.com/apache/spark/pull/57827#discussion_r3808074619
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -1603,6 +1659,74 @@ object VariantExplode {
case _ => Nil
}
}
+
+ private case class ExplodeEntry(path: String, pos: Int, key: UTF8String,
value: Variant)
+
+ def variantExplodeRecursive(
+ input: VariantVal,
+ isNull: Boolean): Iterable[InternalRow] = {
+ if (isNull) {
+ return Iterable.empty
+ }
+
+ new Iterable[InternalRow] {
+ override def iterator: Iterator[InternalRow] = {
+ val stack = new ArrayDeque[ExplodeEntry]()
+ pushChildren(new Variant(input.getValue, input.getMetadata), "$",
stack)
+ new Iterator[InternalRow] {
+ override def hasNext: Boolean = !stack.isEmpty
+
+ override def next(): InternalRow = {
+ val entry = stack.pop()
+ pushChildren(entry.value, entry.path, stack)
+ InternalRow(
+ UTF8String.fromString(entry.path),
+ entry.pos,
+ entry.key,
+ new VariantVal(entry.value.getValue, entry.value.getMetadata))
+ }
+ }
+ }
+ }
+ }
+
+ private def pushChildren(
+ v: Variant,
+ parentPath: String,
+ stack: ArrayDeque[ExplodeEntry]): Unit = {
+ v.getType match {
+ case Type.OBJECT =>
+ for (i <- v.objectSize() - 1 to 0 by -1) {
+ val field = v.getFieldAtIndex(i)
+ stack.push(ExplodeEntry(
+ appendObjectPath(parentPath, field.key),
+ i,
+ UTF8String.fromString(field.key),
+ field.value))
+ }
+ case Type.ARRAY =>
+ for (i <- v.arraySize() - 1 to 0 by -1) {
+ stack.push(ExplodeEntry(
+ s"$parentPath[$i]",
+ i,
+ null,
+ v.getElementAtIndex(i)))
+ }
+ case _ =>
+ }
+ }
+
+ private def appendObjectPath(parentPath: String, key: String): String = {
+ if (key.nonEmpty && !key.contains('.') && !key.contains('[')) {
+ s"$parentPath.$key"
+ } else if (!key.contains('"')) {
+ s"""$parentPath["$key"]"""
+ } else if (!key.contains('\'')) {
+ s"$parentPath['$key']"
+ } else {
+ s"""$parentPath["${StringEscapeUtils.escapeJson(key)}"]"""
Review Comment:
It may be difficult to ensure the emitted path is always a valid JSON path
(say, if we expect the user to use the path as the input of `variant_get`),
because the JSON path parsing doesn't have the ability to unescape characters.
But maybe it is not a blocker for this PR.
##########
python/pyspark/sql/connect/tvf.py:
##########
@@ -99,13 +99,17 @@ def sql_keywords(self) -> "DataFrame":
sql_keywords.__doc__ = PySparkTableValuedFunction.sql_keywords.__doc__
- def variant_explode(self, input: "Column") -> "DataFrame":
- return self._fn("variant_explode", input)
+ def variant_explode(self, input: "Column", recursive: bool = False) ->
"DataFrame":
+ from pyspark.sql.connect.functions.builtin import lit
+
+ return self._fn("variant_explode", input, cast("Column",
lit(recursive)))
Review Comment:
Pasting an AI comment, not sure whether this is a concern.
The Connect client now always appends `lit(recursive)`, so
`spark.tvf.variant_explode(col)` sends 2 args on the default path. A pre-4.4
server's functionSignature has one parameter → WRONG_NUM_ARGS. A client-only
upgrade breaks code that never asked for recursion.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -1603,6 +1659,74 @@ object VariantExplode {
case _ => Nil
}
}
+
+ private case class ExplodeEntry(path: String, pos: Int, key: UTF8String,
value: Variant)
+
+ def variantExplodeRecursive(
+ input: VariantVal,
+ isNull: Boolean): Iterable[InternalRow] = {
+ if (isNull) {
+ return Iterable.empty
+ }
+
+ new Iterable[InternalRow] {
+ override def iterator: Iterator[InternalRow] = {
+ val stack = new ArrayDeque[ExplodeEntry]()
+ pushChildren(new Variant(input.getValue, input.getMetadata), "$",
stack)
+ new Iterator[InternalRow] {
+ override def hasNext: Boolean = !stack.isEmpty
+
+ override def next(): InternalRow = {
+ val entry = stack.pop()
+ pushChildren(entry.value, entry.path, stack)
+ InternalRow(
+ UTF8String.fromString(entry.path),
+ entry.pos,
+ entry.key,
+ new VariantVal(entry.value.getValue, entry.value.getMetadata))
+ }
+ }
+ }
+ }
+ }
+
+ private def pushChildren(
+ v: Variant,
+ parentPath: String,
+ stack: ArrayDeque[ExplodeEntry]): Unit = {
+ v.getType match {
+ case Type.OBJECT =>
+ for (i <- v.objectSize() - 1 to 0 by -1) {
+ val field = v.getFieldAtIndex(i)
+ stack.push(ExplodeEntry(
+ appendObjectPath(parentPath, field.key),
+ i,
+ UTF8String.fromString(field.key),
+ field.value))
+ }
+ case Type.ARRAY =>
+ for (i <- v.arraySize() - 1 to 0 by -1) {
+ stack.push(ExplodeEntry(
+ s"$parentPath[$i]",
+ i,
+ null,
+ v.getElementAtIndex(i)))
+ }
+ case _ =>
+ }
+ }
+
+ private def appendObjectPath(parentPath: String, key: String): String = {
+ if (key.nonEmpty && !key.contains('.') && !key.contains('[')) {
+ s"$parentPath.$key"
Review Comment:
The dot-notation branch only rejects `.` and `[`, so keys containing spaces,
`]`, `"`, `$`, `*`, backslashes or control characters are emitted as bare
`$.key`, which is not the valid JSONPath the docs promise.
--
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]
