This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 819fba1b279c [SPARK-58645][SQL] Add `recursive` flag to
`variant_explode`
819fba1b279c is described below
commit 819fba1b279c27163c7a01bbae07d0ddacd394bb
Author: bojana-db <[email protected]>
AuthorDate: Thu Oct 29 16:19:40 2026 +0000
[SPARK-58645][SQL] Add `recursive` flag to `variant_explode`
## What changes were proposed in this pull request?
Adds an optional `recursive` argument to `variant_explode` and
`variant_explode_outer`. When enabled, nested Variant objects and arrays are
recursively exploded and a path column is added to the output.
## Why are the changes needed?
Currently, nested Variant values require repeated lateral calls to
`variant_explode`.
## Does this PR introduce any user-facing change?
Yes. SQL, Scala, and Python APIs now support recursive Variant explosion.
Existing behavior remains unchanged by default.
## How was this patch tested?
Unit tests.
## Was this patch authored or co-authored using generative AI tooling?
Generated-by: GPT-5.6 Sol
Closes #57827 from bojana-db/recursive-explode.
Authored-by: bojana-db <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
project/MimaExcludes.scala | 7 +-
python/pyspark/sql/connect/tvf.py | 20 ++-
python/pyspark/sql/tests/test_tvf.py | 14 ++
python/pyspark/sql/tvf.py | 67 +++++++-
.../org/apache/spark/sql/TableValuedFunction.scala | 32 ++++
.../expressions/variant/variantExpressions.scala | 154 +++++++++++++++--
.../sql/DataFrameTableValuedFunctionsSuite.scala | 14 ++
.../spark/sql/connect/TableValuedFunction.scala | 9 +
.../spark/sql/classic/TableValuedFunction.scala | 9 +
.../variant/named-function-arguments.sql.out | 30 +++-
.../inputs/variant/named-function-arguments.sql | 2 +
.../variant/named-function-arguments.sql.out | 20 +++
.../sql/DataFrameTableValuedFunctionsSuite.scala | 24 +++
.../scala/org/apache/spark/sql/VariantSuite.scala | 190 ++++++++++++++++++++-
14 files changed, 556 insertions(+), 36 deletions(-)
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index f2029f1e97d9..eaaa35939f6a 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -48,7 +48,12 @@ object MimaExcludes {
ProblemFilters.exclude[DirectMissingMethodProblem](
"org.apache.spark.ml.regression.DecisionTreeRegressionModel.predictVariance"),
// BIN BY DataFrame API.
-
ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.sql.Dataset.binBy")
+
ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.sql.Dataset.binBy"),
+ // [SPARK-58645][SQL] Add recursive mode to variant explode table-valued
functions
+ ProblemFilters.exclude[ReversedMissingMethodProblem](
+ "org.apache.spark.sql.TableValuedFunction.variant_explode"),
+ ProblemFilters.exclude[ReversedMissingMethodProblem](
+ "org.apache.spark.sql.TableValuedFunction.variant_explode_outer")
)
// Exclude rules for 4.3.x from 4.2.0 (add 4.3-specific filters below as
needed).
diff --git a/python/pyspark/sql/connect/tvf.py
b/python/pyspark/sql/connect/tvf.py
index 395ade38bdc5..582104b8a76e 100644
--- a/python/pyspark/sql/connect/tvf.py
+++ b/python/pyspark/sql/connect/tvf.py
@@ -14,7 +14,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, cast
from pyspark.errors import PySparkValueError
from pyspark.sql.tvf import TableValuedFunction as PySparkTableValuedFunction
@@ -99,13 +99,23 @@ class TableValuedFunction:
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":
+ if not recursive:
+ return self._fn("variant_explode", input)
+
+ from pyspark.sql.connect.functions.builtin import lit
+
+ return self._fn("variant_explode", input, cast("Column",
lit(recursive)))
variant_explode.__doc__ =
PySparkTableValuedFunction.variant_explode.__doc__
- def variant_explode_outer(self, input: "Column") -> "DataFrame":
- return self._fn("variant_explode_outer", input)
+ def variant_explode_outer(self, input: "Column", recursive: bool = False)
-> "DataFrame":
+ if not recursive:
+ return self._fn("variant_explode_outer", input)
+
+ from pyspark.sql.connect.functions.builtin import lit
+
+ return self._fn("variant_explode_outer", input, cast("Column",
lit(recursive)))
variant_explode_outer.__doc__ =
PySparkTableValuedFunction.variant_explode_outer.__doc__
diff --git a/python/pyspark/sql/tests/test_tvf.py
b/python/pyspark/sql/tests/test_tvf.py
index e0eddba9a41b..c1bd70243467 100644
--- a/python/pyspark/sql/tests/test_tvf.py
+++ b/python/pyspark/sql/tests/test_tvf.py
@@ -589,6 +589,20 @@ class TVFTestsMixin:
expected = self.spark.sql("""SELECT * FROM
variant_explode_outer(parse_json('1'))""")
assertDataFrameEqual(actual=actual, expected=expected)
+ def test_variant_explode_with_recursive(self):
+ nested = sf.parse_json(sf.lit('{"a": [1, {"b": 2}]}'))
+ actual = self.spark.tvf.variant_explode(nested, recursive=True)
+ expected = self.spark.sql(
+ """SELECT * FROM variant_explode(
+ parse_json('{"a": [1, {"b": 2}]}'), true)"""
+ )
+ assertDataFrameEqual(actual=actual, expected=expected)
+
+ def test_variant_explode_outer_with_recursive(self):
+ actual =
self.spark.tvf.variant_explode_outer(sf.parse_json(sf.lit("[]")),
recursive=True)
+ expected = self.spark.sql("""SELECT * FROM
variant_explode_outer(parse_json('[]'), true)""")
+ assertDataFrameEqual(actual=actual, expected=expected)
+
def test_variant_explode_outer_with_lateral_join(self):
with self.temp_view("variant_table"):
variant_table = self.spark.sql("""
diff --git a/python/pyspark/sql/tvf.py b/python/pyspark/sql/tvf.py
index b32f62ab61e5..d702b9a0d816 100644
--- a/python/pyspark/sql/tvf.py
+++ b/python/pyspark/sql/tvf.py
@@ -555,22 +555,28 @@ class TableValuedFunction:
"""
return self._fn("sql_keywords")
- def variant_explode(self, input: Column) -> DataFrame:
+ def variant_explode(self, input: Column, recursive: bool = False) ->
DataFrame:
"""
Separates a variant object/array into multiple rows containing its
fields/elements.
Its result schema is `struct<pos int, key string, value variant>`.
`pos` is the position of
the field/element in its parent object/array, and `value` is the
field/element value.
`key` is the field name when exploding a variant object, or is NULL
when exploding a variant
- array. It ignores any input that is not a variant array/object,
including SQL NULL, variant
- null, and any other variant values.
+ array. When ``recursive`` is true, it also emits all nested
descendants and adds a leading
+ ``path`` column containing each field/element's JSONPath. It ignores
any input that is not
+ a variant array/object, including SQL NULL, variant null, and any
other variant values.
.. versionadded:: 4.0.0
+ .. versionchanged:: 4.4.0
+ Added the ``recursive`` parameter.
+
Parameters
----------
input : :class:`~pyspark.sql.Column`
input column of values to explode.
+ recursive : bool, optional
+ whether to recursively explode nested objects and arrays, by
default False.
Returns
-------
@@ -617,25 +623,51 @@ class TableValuedFunction:
|pos|key|value|
+---+---+-----+
+---+---+-----+
+
+ Example 5: Recursively exploding a nested variant
+
+ >>> value = sf.parse_json(sf.lit('{"a": [1, 2]}'))
+ >>> spark.tvf.variant_explode(value, recursive=True).show()
+ +------+---+----+-----+
+ | path|pos| key|value|
+ +------+---+----+-----+
+ | $.a| 0| a|[1,2]|
+ |$.a[0]| 0|NULL| 1|
+ |$.a[1]| 1|NULL| 2|
+ +------+---+----+-----+
"""
- return self._fn("variant_explode", input)
+ from pyspark.sql.classic.column import _to_java_column
+
+ return DataFrame(
+ self._sparkSession._jsparkSession.tvf().variant_explode(
+ _to_java_column(input), recursive
+ ),
+ self._sparkSession,
+ )
- def variant_explode_outer(self, input: Column) -> DataFrame:
+ def variant_explode_outer(self, input: Column, recursive: bool = False) ->
DataFrame:
"""
Separates a variant object/array into multiple rows containing its
fields/elements.
Its result schema is `struct<pos int, key string, value variant>`.
`pos` is the position of
the field/element in its parent object/array, and `value` is the
field/element value.
`key` is the field name when exploding a variant object, or is NULL
when exploding a variant
- array. Unlike variant_explode, if the given variant is not a variant
array/object, including
- SQL NULL, variant null, and any other variant values, then NULL is
produced.
+ array. When ``recursive`` is true, it also emits all nested
descendants and adds a leading
+ ``path`` column containing each field/element's JSONPath. Unlike
variant_explode, if the
+ given variant is not a variant array/object, including SQL NULL,
variant null, and any other
+ variant values, then NULL is produced.
.. versionadded:: 4.0.0
+ .. versionchanged:: 4.4.0
+ Added the ``recursive`` parameter.
+
Parameters
----------
input : :class:`~pyspark.sql.Column`
input column of values to explode.
+ recursive : bool, optional
+ whether to recursively explode nested objects and arrays, by
default False.
Returns
-------
@@ -670,8 +702,27 @@ class TableValuedFunction:
+----+----+-----+
|NULL|NULL| NULL|
+----+----+-----+
+
+ Recursively exploding a nested variant
+
+ >>> value = sf.parse_json(sf.lit('{"a": [1, 2]}'))
+ >>> spark.tvf.variant_explode_outer(value, recursive=True).show()
+ +------+---+----+-----+
+ | path|pos| key|value|
+ +------+---+----+-----+
+ | $.a| 0| a|[1,2]|
+ |$.a[0]| 0|NULL| 1|
+ |$.a[1]| 1|NULL| 2|
+ +------+---+----+-----+
"""
- return self._fn("variant_explode_outer", input)
+ from pyspark.sql.classic.column import _to_java_column
+
+ return DataFrame(
+ self._sparkSession._jsparkSession.tvf().variant_explode_outer(
+ _to_java_column(input), recursive
+ ),
+ self._sparkSession,
+ )
def python_worker_logs(self) -> DataFrame:
"""
diff --git
a/sql/api/src/main/scala/org/apache/spark/sql/TableValuedFunction.scala
b/sql/api/src/main/scala/org/apache/spark/sql/TableValuedFunction.scala
index 56f2f75c2c19..a8381a02fd52 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/TableValuedFunction.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/TableValuedFunction.scala
@@ -169,6 +169,22 @@ abstract class TableValuedFunction {
*/
def variant_explode(input: Column): Dataset[Row]
+ /**
+ * Separates a variant object/array into multiple rows, including nested
fields/elements when
+ * `recursive` is true. When false, the result schema is
+ * `struct<pos int, key string, value variant>`; when true, the result
schema is
+ * `struct<path string, pos int, key string, value variant>`.
+ *
+ * `pos` is the position within the parent, `key` is the object field name
or NULL for arrays,
+ * and `value` is the field/element value. In recursive mode, `path` is its
JSONPath. It ignores
+ * any input that is not a variant array/object, including SQL NULL, variant
null, and any other
+ * variant values.
+ *
+ * @group variant_funcs
+ * @since 4.4.0
+ */
+ def variant_explode(input: Column, recursive: Boolean): Dataset[Row]
+
/**
* Separates a variant object/array into multiple rows containing its
fields/elements. Its
* result schema is `struct<pos int, key string, value variant>`.
`pos` is the position of
@@ -182,6 +198,22 @@ abstract class TableValuedFunction {
*/
def variant_explode_outer(input: Column): Dataset[Row]
+ /**
+ * Separates a variant object/array into multiple rows, including nested
fields/elements when
+ * `recursive` is true. When false, the result schema is
+ * `struct<pos int, key string, value variant>`; when true, the result
schema is
+ * `struct<path string, pos int, key string, value variant>`.
+ *
+ * `pos` is the position within the parent, `key` is the object field name
or NULL for arrays,
+ * and `value` is the field/element value. In recursive mode, `path` is its
JSONPath. Unlike
+ * variant_explode, if the given variant is not a variant array/object,
including SQL NULL,
+ * variant null, and any other variant values, then NULL is produced.
+ *
+ * @group variant_funcs
+ * @since 4.4.0
+ */
+ def variant_explode_outer(input: Column, recursive: Boolean): Dataset[Row]
+
/**
* Returns a `DataFrame` of logs collected from Python workers.
*
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
index f4aeacc4ddb7..1b2aa10eccf8 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala
@@ -18,6 +18,7 @@
package org.apache.spark.sql.catalyst.expressions.variant
import java.time.ZoneId
+import java.util.ArrayDeque
import scala.util.matching.Regex
import scala.util.parsing.combinator.RegexParsers
@@ -1934,33 +1935,76 @@ object VariantStripNullsExpressionBuilder extends
ExpressionBuilder {
}
}
-case class VariantExplode(child: Expression) extends UnaryExpression with
Generator
- with ExpectsInputTypes {
- override def inputTypes: Seq[AbstractDataType] = Seq(VariantType)
+case class VariantExplode(child: Expression, recursive: Expression)
+ extends BinaryExpression with Generator
+ with ExpectsInputTypes
+ with QueryErrorsBase {
+ override def left: Expression = child
+ override def right: Expression = recursive
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(VariantType,
BooleanType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val check = super.checkInputDataTypes()
+ if (check.isFailure) {
+ check
+ } else if (!recursive.foldable) {
+ DataTypeMismatch(
+ errorSubClass = "NON_FOLDABLE_INPUT",
+ messageParameters = Map(
+ "inputName" -> toSQLId("recursive"),
+ "inputType" -> toSQLType(recursive.dataType),
+ "inputExpr" -> toSQLExpr(recursive)))
+ } else if (recursive.eval() == null) {
+ DataTypeMismatch(
+ errorSubClass = "UNEXPECTED_NULL",
+ messageParameters = Map("exprName" -> "recursive"))
+ } else {
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+
+ private lazy val isRecursive = recursive.foldable && (recursive.eval() match
{
+ case value: Boolean => value
+ case _ => false
+ })
override def prettyName: String = "variant_explode"
- override protected def withNewChildInternal(newChild: Expression):
VariantExplode =
- copy(child = newChild)
+ override protected def withNewChildrenInternal(
+ newLeft: Expression,
+ newRight: Expression): VariantExplode =
+ copy(child = newLeft, recursive = newRight)
override def eval(input: InternalRow): IterableOnce[InternalRow] = {
val inputVariant = child.eval(input).asInstanceOf[VariantVal]
- VariantExplode.variantExplode(inputVariant, inputVariant == null)
+ if (isRecursive) {
+ VariantExplode.variantExplodeRecursive(inputVariant, inputVariant ==
null)
+ } else {
+ VariantExplode.variantExplode(inputVariant, inputVariant == null)
+ }
}
override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
val childCode = child.genCode(ctx)
val cls = classOf[VariantExplode].getName
+ val method = if (isRecursive) "variantExplodeRecursive" else
"variantExplode"
val code = code"""
${childCode.code}
- scala.collection.Seq<InternalRow> ${ev.value} = $cls.variantExplode(
+ scala.collection.IterableOnce<InternalRow> ${ev.value} = $cls.$method(
${childCode.value}, ${childCode.isNull});
"""
ev.copy(code = code, isNull = FalseLiteral)
}
override def elementSchema: StructType = {
- new StructType()
+ val schema = new StructType()
+ val schemaWithPath = if (isRecursive) {
+ schema.add("path", StringType, nullable = false)
+ } else {
+ schema
+ }
+ schemaWithPath
.add("pos", IntegerType, nullable = false)
.add("key", StringType, nullable = true)
.add("value", VariantType, nullable = false)
@@ -1968,17 +2012,19 @@ case class VariantExplode(child: Expression) extends
UnaryExpression with Genera
}
trait VariantExplodeGeneratorBuilderBase extends GeneratorBuilder {
- override def functionSignature: Option[FunctionSignature] =
- Some(FunctionSignature(Seq(InputParameter("input"))))
+ override def functionSignature: Option[FunctionSignature] =
Some(FunctionSignature(Seq(
+ InputParameter("input"),
+ InputParameter("recursive", Some(Literal.create(false, BooleanType))))))
+
override def buildGenerator(funcName: String, expressions: Seq[Expression]):
Generator = {
- assert(expressions.size == 1)
- VariantExplode(expressions(0))
+ assert(expressions.size == 2)
+ VariantExplode(expressions(0), expressions(1))
}
}
// scalastyle:off line.size.limit line.contains.tab
@ExpressionDescription(
- usage = "_FUNC_(expr) - It separates a variant object/array into multiple
rows containing its fields/elements. Its result schema is `struct<pos int, key
string, value variant>`. `pos` is the position of the field/element in its
parent object/array, and `value` is the field/element value. `key` is the field
name when exploding a variant object, or is NULL when exploding a variant
array. It ignores any input that is not a variant array/object, including SQL
NULL, variant null, and any ot [...]
+ usage = "_FUNC_(expr[, recursive]) - It separates a variant object/array
into multiple rows containing its fields/elements. Its result schema is
`struct<pos int, key string, value variant>`. When `recursive` is true, it also
emits all nested descendants and adds a leading `path string` column containing
each field/element's JSONPath. `pos` is the position of the field/element in
its parent object/array, and `value` is the field/element value. `key` is the
field name when exploding a va [...]
examples = """
Examples:
> SELECT * from _FUNC_(parse_json('["hello", "world"]'));
@@ -1987,6 +2033,10 @@ trait VariantExplodeGeneratorBuilderBase extends
GeneratorBuilder {
> SELECT * from _FUNC_(input => parse_json('{"a": true, "b": 3.14}'));
0 a true
1 b 3.14
+ > SELECT * from _FUNC_(parse_json('{"a": [1, 2]}'), true);
+ $.a 0 a [1,2]
+ $.a[0] 0 NULL 1
+ $.a[1] 1 NULL 2
""",
since = "4.0.0",
group = "variant_funcs")
@@ -1997,7 +2047,7 @@ object VariantExplodeGeneratorBuilder extends
VariantExplodeGeneratorBuilderBase
// scalastyle:off line.size.limit line.contains.tab
@ExpressionDescription(
- usage = "_FUNC_(expr) - It separates a variant object/array into multiple
rows containing its fields/elements. Its result schema is `struct<pos int, key
string, value variant>`. `pos` is the position of the field/element in its
parent object/array, and `value` is the field/element value. `key` is the field
name when exploding a variant object, or is NULL when exploding a variant
array. It ignores any input that is not a variant array/object, including SQL
NULL, variant null, and any ot [...]
+ usage = "_FUNC_(expr[, recursive]) - It separates a variant object/array
into multiple rows containing its fields/elements. Its result schema is
`struct<pos int, key string, value variant>`. When `recursive` is true, it also
emits all nested descendants and adds a leading `path string` column containing
each field/element's JSONPath. `pos` is the position of the field/element in
its parent object/array, and `value` is the field/element value. `key` is the
field name when exploding a va [...]
examples = """
Examples:
> SELECT * from _FUNC_(parse_json('["hello", "world"]'));
@@ -2006,6 +2056,10 @@ object VariantExplodeGeneratorBuilder extends
VariantExplodeGeneratorBuilderBase
> SELECT * from _FUNC_(input => parse_json('{"a": true, "b": 3.14}'));
0 a true
1 b 3.14
+ > SELECT * from _FUNC_(parse_json('{"a": [1, 2]}'), true);
+ $.a 0 a [1,2]
+ $.a[0] 0 NULL 1
+ $.a[1] 1 NULL 2
""",
since = "4.0.0",
group = "variant_funcs")
@@ -2046,6 +2100,78 @@ object VariantExplode {
case _ => Nil
}
}
+
+ private case class ExplodeEntry(path: String, pos: Int, key: UTF8String,
value: Variant)
+
+ def variantExplodeRecursive(
+ input: VariantVal,
+ isNull: Boolean): Iterator[InternalRow] = {
+ if (isNull) {
+ return Iterator.empty
+ }
+
+ 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 _ =>
+ }
+ }
+
+ // Use dot notation for dot-safe keys and single-quoted bracket notation for
other keys.
+ // Emitted paths target the default escaped-path parser, regardless of the
legacy parser setting.
+ private def appendObjectPath(parentPath: String, key: String): String = {
+ if (isDotSafeKey(key)) {
+ s"$parentPath.$key"
+ } else {
+ s"$parentPath${VariantPathParser.quoteKey(key, '\'')}"
+ }
+ }
+
+ // Dot-safe keys are non-empty and identifier-like: an ASCII letter, `_`, or
non-ASCII character
+ // first, then those or ASCII digits. All other keys use bracket notation.
+ private def isDotSafeKey(key: String): Boolean = {
+ def isStart(c: Char): Boolean =
+ (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || c == '_' || c >= 0x80
+ def isPart(c: Char): Boolean = isStart(c) || (c >= '0' && c <= '9')
+ key.nonEmpty && isStart(key.charAt(0)) &&
+ (1 until key.length).forall(i => isPart(key.charAt(i)))
+ }
}
@ExpressionDescription(
diff --git
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
index e619279940d2..c9f5e0272f4a 100644
---
a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
+++
b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
@@ -524,6 +524,20 @@ class DataFrameTableValuedFunctionsSuite extends QueryTest
with RemoteSparkSessi
}
}
+ test("variant_explode - recursive") {
+ val nested = parse_json(lit("""{"a": [1, {"b": 2}]}"""))
+ val actual = spark.tvf.variant_explode(nested, recursive = true)
+ val expected =
+ spark.sql("""SELECT * FROM variant_explode(parse_json('{"a": [1, {"b":
2}]}'), true)""")
+ checkAnswer(actual, expected)
+ }
+
+ test("variant_explode_outer - recursive") {
+ val actual = spark.tvf.variant_explode_outer(parse_json(lit("[]")),
recursive = true)
+ val expected = spark.sql("SELECT * FROM
variant_explode_outer(parse_json('[]'), true)")
+ checkAnswer(actual, expected)
+ }
+
test("explode with udf") {
Seq("NO_CODEGEN", "CODEGEN_ONLY").foreach { codegenMode =>
withSQLConf("spark.sql.codegen.factoryMode" -> codegenMode) {
diff --git
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/TableValuedFunction.scala
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/TableValuedFunction.scala
index e773f66cd6d0..e28584c42718 100644
---
a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/TableValuedFunction.scala
+++
b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/TableValuedFunction.scala
@@ -20,6 +20,7 @@ import scala.jdk.CollectionConverters._
import org.apache.spark.sql
import org.apache.spark.sql.{Column, Row}
+import org.apache.spark.sql.functions.lit
class TableValuedFunction(sparkSession: SparkSession) extends
sql.TableValuedFunction {
@@ -99,10 +100,18 @@ class TableValuedFunction(sparkSession: SparkSession)
extends sql.TableValuedFun
override def variant_explode(input: Column): Dataset[Row] =
fn("variant_explode", Seq(input))
+ /** @inheritdoc */
+ override def variant_explode(input: Column, recursive: Boolean):
Dataset[Row] =
+ fn("variant_explode", Seq(input, lit(recursive)))
+
/** @inheritdoc */
override def variant_explode_outer(input: Column): Dataset[Row] =
fn("variant_explode_outer", Seq(input))
+ /** @inheritdoc */
+ override def variant_explode_outer(input: Column, recursive: Boolean):
Dataset[Row] =
+ fn("variant_explode_outer", Seq(input, lit(recursive)))
+
/** @inheritdoc */
override def python_worker_logs(): Dataset[Row] =
fn("python_worker_logs", Seq.empty)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/classic/TableValuedFunction.scala
b/sql/core/src/main/scala/org/apache/spark/sql/classic/TableValuedFunction.scala
index 839faa85aed4..84f72406efdd 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/classic/TableValuedFunction.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/classic/TableValuedFunction.scala
@@ -19,6 +19,7 @@ package org.apache.spark.sql.classic
import org.apache.spark.sql
import org.apache.spark.sql.{Column, Row}
import org.apache.spark.sql.catalyst.analysis.UnresolvedTableValuedFunction
+import org.apache.spark.sql.functions.lit
class TableValuedFunction(sparkSession: SparkSession)
extends sql.TableValuedFunction {
@@ -94,10 +95,18 @@ class TableValuedFunction(sparkSession: SparkSession)
override def variant_explode(input: Column): Dataset[Row] =
fn("variant_explode", Seq(input))
+ /** @inheritdoc */
+ override def variant_explode(input: Column, recursive: Boolean):
Dataset[Row] =
+ fn("variant_explode", Seq(input, lit(recursive)))
+
/** @inheritdoc */
override def variant_explode_outer(input: Column): Dataset[Row] =
fn("variant_explode_outer", Seq(input))
+ /** @inheritdoc */
+ override def variant_explode_outer(input: Column, recursive: Boolean):
Dataset[Row] =
+ fn("variant_explode_outer", Seq(input, lit(recursive)))
+
/** @inheritdoc */
override def python_worker_logs(): Dataset[Row] =
fn("python_worker_logs", Seq.empty)
diff --git
a/sql/core/src/test/resources/sql-tests/analyzer-results/variant/named-function-arguments.sql.out
b/sql/core/src/test/resources/sql-tests/analyzer-results/variant/named-function-arguments.sql.out
index 571ec68048cc..72fa804b53e6 100644
---
a/sql/core/src/test/resources/sql-tests/analyzer-results/variant/named-function-arguments.sql.out
+++
b/sql/core/src/test/resources/sql-tests/analyzer-results/variant/named-function-arguments.sql.out
@@ -3,7 +3,7 @@
SELECT * FROM variant_explode(input => parse_json('["hello", "world"]'))
-- !query analysis
Project [pos#x, key#x, value#x]
-+- Generate variant_explode(parse_json(["hello", "world"], true)), false,
[pos#x, key#x, value#x]
++- Generate variant_explode(parse_json(["hello", "world"], true), false),
false, [pos#x, key#x, value#x]
+- OneRowRelation
@@ -11,7 +11,23 @@ Project [pos#x, key#x, value#x]
SELECT * FROM variant_explode_outer(input => parse_json('{"a": true, "b":
3.14}'))
-- !query analysis
Project [pos#x, key#x, value#x]
-+- Generate variant_explode(parse_json({"a": true, "b": 3.14}, true)), true,
[pos#x, key#x, value#x]
++- Generate variant_explode(parse_json({"a": true, "b": 3.14}, true), false),
true, [pos#x, key#x, value#x]
+ +- OneRowRelation
+
+
+-- !query
+SELECT * FROM variant_explode(input => parse_json('{"a": [1, 2]}'), recursive
=> true)
+-- !query analysis
+Project [path#x, pos#x, key#x, value#x]
++- Generate variant_explode(parse_json({"a": [1, 2]}, true), true), false,
[path#x, pos#x, key#x, value#x]
+ +- OneRowRelation
+
+
+-- !query
+SELECT * FROM variant_explode_outer(recursive => true, input =>
parse_json('{"a": [1, 2]}'))
+-- !query analysis
+Project [path#x, pos#x, key#x, value#x]
++- Generate variant_explode(parse_json({"a": [1, 2]}, true), true), true,
[path#x, pos#x, key#x, value#x]
+- OneRowRelation
@@ -20,9 +36,9 @@ SELECT * FROM variant_explode(parse_json('["hello",
"world"]')), variant_explode
-- !query analysis
Project [pos#x, key#x, value#x, pos#x, key#x, value#x]
+- Join Inner
- :- Generate variant_explode(parse_json(["hello", "world"], true)), false,
[pos#x, key#x, value#x]
+ :- Generate variant_explode(parse_json(["hello", "world"], true), false),
false, [pos#x, key#x, value#x]
: +- OneRowRelation
- +- Generate variant_explode(parse_json({"a": true, "b": 3.14}, true)),
false, [pos#x, key#x, value#x]
+ +- Generate variant_explode(parse_json({"a": true, "b": 3.14}, true),
false), false, [pos#x, key#x, value#x]
+- OneRowRelation
@@ -31,10 +47,10 @@ SELECT * FROM variant_explode(parse_json('{"a": ["hello",
"world"], "b": {"x": t
-- !query analysis
Project [pos#x, key#x, value#x, pos#x, key#x, value#x]
+- LateralJoin lateral-subquery#x [value#x], Inner
- : +- Generate variant_explode(outer(value#x)), false, [pos#x, key#x,
value#x]
+ : +- Generate variant_explode(outer(value#x), false), false, [pos#x,
key#x, value#x]
: +- OneRowRelation
+- SubqueryAlias t
- +- Generate variant_explode(parse_json({"a": ["hello", "world"], "b":
{"x": true, "y": 3.14}}, true)), false, [pos#x, key#x, value#x]
+ +- Generate variant_explode(parse_json({"a": ["hello", "world"], "b":
{"x": true, "y": 3.14}}, true), false), false, [pos#x, key#x, value#x]
+- OneRowRelation
@@ -44,5 +60,5 @@ SELECT num, key, val, 'Spark' FROM
variant_explode(parse_json('["hello", "world"
Project [num#x, key#x, val#x, Spark AS Spark#x]
+- SubqueryAlias t
+- Project [pos#x AS num#x, key#x AS key#x, value#x AS val#x]
- +- Generate variant_explode(parse_json(["hello", "world"], true)),
false, [pos#x, key#x, value#x]
+ +- Generate variant_explode(parse_json(["hello", "world"], true),
false), false, [pos#x, key#x, value#x]
+- OneRowRelation
diff --git
a/sql/core/src/test/resources/sql-tests/inputs/variant/named-function-arguments.sql
b/sql/core/src/test/resources/sql-tests/inputs/variant/named-function-arguments.sql
index e63e68169a81..552046ef17b0 100644
---
a/sql/core/src/test/resources/sql-tests/inputs/variant/named-function-arguments.sql
+++
b/sql/core/src/test/resources/sql-tests/inputs/variant/named-function-arguments.sql
@@ -1,6 +1,8 @@
-- Test for tabled value functions variant_explode and variant_explode_outer
SELECT * FROM variant_explode(input => parse_json('["hello", "world"]'));
SELECT * FROM variant_explode_outer(input => parse_json('{"a": true, "b":
3.14}'));
+SELECT * FROM variant_explode(input => parse_json('{"a": [1, 2]}'), recursive
=> true);
+SELECT * FROM variant_explode_outer(recursive => true, input =>
parse_json('{"a": [1, 2]}'));
SELECT * FROM variant_explode(parse_json('["hello", "world"]')),
variant_explode(parse_json('{"a": true, "b": 3.14}'));
SELECT * FROM variant_explode(parse_json('{"a": ["hello", "world"], "b": {"x":
true, "y": 3.14}}')) AS t, LATERAL variant_explode(t.value);
SELECT num, key, val, 'Spark' FROM variant_explode(parse_json('["hello",
"world"]')) AS t(num, key, val);
diff --git
a/sql/core/src/test/resources/sql-tests/results/variant/named-function-arguments.sql.out
b/sql/core/src/test/resources/sql-tests/results/variant/named-function-arguments.sql.out
index 10a439831fbd..df1c2ce7d44d 100644
---
a/sql/core/src/test/resources/sql-tests/results/variant/named-function-arguments.sql.out
+++
b/sql/core/src/test/resources/sql-tests/results/variant/named-function-arguments.sql.out
@@ -17,6 +17,26 @@ struct<pos:int,key:string,value:variant>
1 b 3.14
+-- !query
+SELECT * FROM variant_explode(input => parse_json('{"a": [1, 2]}'), recursive
=> true)
+-- !query schema
+struct<path:string,pos:int,key:string,value:variant>
+-- !query output
+$.a 0 a [1,2]
+$.a[0] 0 NULL 1
+$.a[1] 1 NULL 2
+
+
+-- !query
+SELECT * FROM variant_explode_outer(recursive => true, input =>
parse_json('{"a": [1, 2]}'))
+-- !query schema
+struct<path:string,pos:int,key:string,value:variant>
+-- !query output
+$.a 0 a [1,2]
+$.a[0] 0 NULL 1
+$.a[1] 1 NULL 2
+
+
-- !query
SELECT * FROM variant_explode(parse_json('["hello", "world"]')),
variant_explode(parse_json('{"a": true, "b": 3.14}'))
-- !query schema
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
index a28cf3fc97d5..aa1e9240e01d 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameTableValuedFunctionsSuite.scala
@@ -528,6 +528,30 @@ class DataFrameTableValuedFunctionsSuite extends
SharedSparkSession {
}
}
+ test("variant_explode - recursive") {
+ val nested = parse_json(lit("""{"a": [1, {"b": 2}]}"""))
+ val actual = spark.tvf.variant_explode(nested, recursive = true)
+ val expected = spark.sql(
+ """SELECT * FROM variant_explode(parse_json('{"a": [1, {"b": 2}]}'),
true)""")
+ checkAnswer(actual, expected)
+
+ val nonRecursive = spark.tvf.variant_explode(nested, recursive = false)
+ assert(nonRecursive.schema.fieldNames.toSeq == Seq("pos", "key", "value"))
+ checkAnswer(nonRecursive, spark.tvf.variant_explode(nested))
+ }
+
+ test("variant_explode_outer - recursive") {
+ val actual = spark.tvf.variant_explode_outer(parse_json(lit("[]")),
recursive = true)
+ val expected = spark.sql(
+ "SELECT * FROM variant_explode_outer(parse_json('[]'), true)")
+ checkAnswer(actual, expected)
+
+ val nested = parse_json(lit("""{"a": [1, {"b": 2}]}"""))
+ val nonRecursive = spark.tvf.variant_explode_outer(nested, recursive =
false)
+ assert(nonRecursive.schema.fieldNames.toSeq == Seq("pos", "key", "value"))
+ checkAnswer(nonRecursive, spark.tvf.variant_explode_outer(nested))
+ }
+
test("explode with udf") {
Seq("NO_CODEGEN", "CODEGEN_ONLY").foreach { codegenMode =>
withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
index e11d396542ad..d4945f8b47c9 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala
@@ -27,7 +27,7 @@ import scala.util.Random
import org.apache.spark.SparkRuntimeException
import org.apache.spark.sql.catalyst.expressions.{CodegenObjectFactoryMode,
ExpressionEvalHelper, Literal}
-import
org.apache.spark.sql.catalyst.expressions.variant.{VariantExpressionEvalUtils,
VariantGet}
+import org.apache.spark.sql.catalyst.expressions.variant.{VariantExplode,
VariantExpressionEvalUtils, VariantGet}
import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, DateTimeUtils,
GenericArrayData}
import org.apache.spark.sql.errors.QueryExecutionErrors.toSQLId
import org.apache.spark.sql.execution.WholeStageCodegenExec
@@ -1226,6 +1226,15 @@ class VariantSuite extends SharedSparkSession with
ExpressionEvalHelper {
)
}
+ test("recursive variant_explode reuses its iterator") {
+ val input =
VariantExpressionEvalUtils.parseJson(UTF8String.fromString("[1,2]"))
+ val output = VariantExplode.variantExplodeRecursive(input, isNull = false)
+ val iterator = output.iterator
+ assert(!iterator.isEmpty)
+ assert(output.iterator eq iterator)
+ assert(iterator.size == 2)
+ }
+
test("variant_explode") {
def check(input: String, expected: Seq[Row]): Unit = {
withView("v") {
@@ -1246,6 +1255,185 @@ class VariantSuite extends SharedSparkSession with
ExpressionEvalHelper {
check("""{"a": [1, 2, 3], "b": true}""", Seq(Row(0, "a", "[1,2,3]"),
Row(1, "b", "true")))
check("""[null, "hello", {}]""",
Seq(Row(0, null, "null"), Row(1, null, "\"hello\""), Row(2, null,
"{}")))
+
+ val nonRecursive = sql(
+ """SELECT * FROM variant_explode(parse_json('{"a": [1]}'), false)""")
+ assert(nonRecursive.schema.fieldNames.toSeq == Seq("pos", "key",
"value"))
+ checkAnswer(
+ nonRecursive.selectExpr("pos", "key", "to_json(value)"),
+ Seq(Row(0, "a", "[1]")))
+
+ val recursiveExpected = Seq(
+ Row("$.a", 0, "a", """[1,{"b":2}]"""),
+ Row("$.a[0]", 0, null, "1"),
+ Row("$.a[1]", 1, null, """{"b":2}"""),
+ Row("$.a[1].b", 0, "b", "2"),
+ Row("$.c", 1, "c", """{"d":[3]}"""),
+ Row("$.c.d", 0, "d", "[3]"),
+ Row("$.c.d[0]", 0, null, "3"))
+ Seq("variant_explode", "variant_explode_outer").foreach { function =>
+ checkAnswer(
+ sql(s"""
+ |SELECT path, pos, key, to_json(value)
+ |FROM $function(
+ | parse_json('{"a": [1, {"b": 2}], "c": {"d": [3]}}'),
+ | recursive => true)
+ |""".stripMargin),
+ recursiveExpected)
+ }
+
+ checkAnswer(
+ sql("""
+ |SELECT path, pos, key, to_json(value)
+ |FROM variant_explode(parse_json('[{"a": 1}, 2]'), true)
+ |""".stripMargin),
+ Seq(
+ Row("$[0]", 0, null, """{"a":1}"""),
+ Row("$[0].a", 0, "a", "1"),
+ Row("$[1]", 1, null, "2")))
+
+ // Keys that are not dot-safe use bracket notation.
+ checkAnswer(
+ sql("""
+ |SELECT path, key, to_json(value)
+ |FROM variant_explode(to_variant_object(
+ | map('a.b', 1, 'a[b', 2, 'a."b', 3, '', 4)), true)
+ |""".stripMargin),
+ Seq(
+ Row("""$['a.b']""", "a.b", "1"),
+ Row("""$['a[b']""", "a[b", "2"),
+ Row("""$['a."b']""", """a."b""", "3"),
+ Row("""$['']""", "", "4")))
+
+ // Bracketed keys escape `\` and `'`; `"` is left literal.
+ checkAnswer(
+ sql("""SELECT path, key FROM variant_explode(
+ | to_variant_object(map('a''b', 1)), true)""".stripMargin),
+ Seq(Row("""$['a\'b']""", "a'b")))
+ checkAnswer(
+ sql("""SELECT path, key FROM variant_explode(
+ | to_variant_object(map('a"''b', 1)), true)""".stripMargin),
+ Seq(Row("""$['a"\'b']""", """a"'b""")))
+ checkAnswer(
+ sql("""SELECT path, key FROM variant_explode(
+ | to_variant_object(map('a\\b', 1)), true)""".stripMargin),
+ Seq(Row("""$['a\\b']""", """a\b""")))
+ checkAnswer(
+ sql("""SELECT path, key FROM variant_explode(
+ | to_variant_object(map('0', 1)), true)""".stripMargin),
+ Seq(Row("""$['0']""", "0")))
+
+ // Short control escapes (`\b \t \n \f \r`). The keys hold real
control chars, injected
+ // into the map so no SQL/JSON escape layer intervenes.
+ val (bsp, tabc, lfc, ffc, crc) = ("\b", "\t", "\n", "\f", "\r")
+ checkAnswer(
+ sql(s"SELECT path FROM variant_explode(to_variant_object(map(" +
+ s"'a${bsp}b', 1, 'a${tabc}c', 2, 'a${lfc}d', 3, 'a${ffc}e', 4,
'a${crc}g', 5)), true)"),
+ Seq(
+ Row("""$['a\bb']"""),
+ Row("""$['a\tc']"""),
+ Row("""$['a\nd']"""),
+ Row("""$['a\fe']"""),
+ Row("""$['a\rg']""")))
+
+ // Other controls use `\uXXXX` (vertical tab U+000B, injected as a
real char); the expected
+ // path is split around `\u` so Scala leaves it literal.
+ val vtab = 11.toChar
+ checkAnswer(
+ sql(s"SELECT path FROM
variant_explode(to_variant_object(map('a${vtab}b', 1)), true)"),
+ Seq(Row("$['a\\" + "u000bb']")))
+
+ withView("variant_explode_round_trip") {
+ Seq(Map(
+ "a'\"b" -> Seq(1),
+ "a\\nb" -> Seq(2),
+ "a\b\f\n\r\tb" -> Seq(3),
+ s"a${vtab}b" -> Seq(4),
+ "" -> Seq(5),
+ "a.b" -> Seq(6))).toDF("input")
+ .selectExpr("to_variant_object(input) AS v")
+ .createOrReplaceTempView("variant_explode_round_trip")
+ Seq("variant_explode", "variant_explode_outer").foreach { function =>
+ checkAnswer(
+ sql(s"""
+ |SELECT to_json(value), to_json(variant_get(v, path,
'variant'))
+ |FROM variant_explode_round_trip, LATERAL $function(v, true)
+ |""".stripMargin),
+ (1 to 6).flatMap { value =>
+ Seq(Row(s"[$value]", s"[$value]"), Row(value.toString,
value.toString))
+ })
+ }
+ }
+
+ // scalastyle:off nonascii
+ checkAnswer(
+ sql("""
+ |SELECT path, key, to_json(value)
+ |FROM variant_explode(parse_json('{"你好": 5}'), true)
+ |""".stripMargin),
+ Seq(Row("$.你好", "你好", "5")))
+
+ checkAnswer(
+ sql("""
+ |SELECT path, key, to_json(value)
+ |FROM variant_explode(parse_json('{"键.名": 5}'), true)
+ |""".stripMargin),
+ Seq(Row("""$['键.名']""", "键.名", "5")))
+ // scalastyle:on nonascii
+
+ // A single path can mix representations per segment: dot for dot-safe
keys, bracket for
+ // the rest, and `[i]` for array indices.
+ checkAnswer(
+ sql("""
+ |SELECT path, to_json(value)
+ |FROM variant_explode(parse_json('{"a": {"b c": [{"d.e": 1}]}}'),
true)
+ |""".stripMargin),
+ Seq(
+ Row("$.a", """{"b c":[{"d.e":1}]}"""),
+ Row("""$.a['b c']""", """[{"d.e":1}]"""),
+ Row("""$.a['b c'][0]""", """{"d.e":1}"""),
+ Row("""$.a['b c'][0]['d.e']""", "1")))
+
+ checkAnswer(
+ sql("""
+ |SELECT path, to_json(value)
+ |FROM variant_explode(parse_json('{"a": {"b''c": {"d e": [1]}}}'),
true)
+ |""".stripMargin),
+ Seq(
+ Row("$.a", """{"b'c":{"d e":[1]}}"""),
+ Row("""$.a['b\'c']""", """{"d e":[1]}"""),
+ Row("""$.a['b\'c']['d e']""", "[1]"),
+ Row("""$.a['b\'c']['d e'][0]""", "1")))
+
+ checkAnswer(
+ sql("SELECT * FROM variant_explode(parse_json('[]'), true)"),
+ Nil)
+ Seq(
+ "parse_json('[]')",
+ "CAST(NULL AS VARIANT)",
+ "parse_json('null')",
+ "parse_json('1')").foreach { input =>
+ checkAnswer(
+ sql(s"SELECT * FROM variant_explode_outer($input, true)"),
+ Seq(Row(null, null, null, null)))
+ }
+
+ val nonFoldable = intercept[AnalysisException] {
+ sql("""
+ |SELECT *
+ |FROM VALUES (true) AS t(recursive),
+ |LATERAL variant_explode(parse_json('[1]'), recursive)
+ |""".stripMargin)
+ }
+ assert(nonFoldable.getCondition ==
"DATATYPE_MISMATCH.NON_FOLDABLE_INPUT")
+
+ val nullRecursive = intercept[AnalysisException] {
+ sql("""
+ |SELECT *
+ |FROM variant_explode(parse_json('[1]'), CAST(NULL AS BOOLEAN))
+ |""".stripMargin)
+ }
+ assert(nullRecursive.getCondition ==
"DATATYPE_MISMATCH.UNEXPECTED_NULL")
}
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
