You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ur...@apache.org on 2026/09/09 14:07:05 UTC
(spark) branch branch-4.x updated: [SPARK-59130][SQL] Add `variant_pick` expression
This is an automated email from the ASF dual-hosted git repository.
uros-b pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new d8dd0bec59c0 [SPARK-59130][SQL] Add `variant_pick` expression
d8dd0bec59c0 is described below
commit d8dd0bec59c0b45e4e61f34c40b317ba697549bd
Author: bojana-db <bo...@databricks.com>
AuthorDate: Wed Sep 9 14:06:40 2026 +0000
[SPARK-59130][SQL] Add `variant_pick` expression
### What changes were proposed in this pull request?
Adds the SQL function `variant_pick(v, path1[, path2, …])`, which keeps only the fields or array elements of a Variant value at one or more JSONPath locations (a projection), preserving their enclosing structure.
Details:
- Object fields keep their enclosing structure; array elements are compacted into a new array in original order;
- Multiple paths are applied as a set (order-independent, deduplicated), and a broader path subsumes a narrower one;
- `NULL` paths are skipped; a `NULL` variant input returns `NULL`;
- paths that don't match (missing key, out-of-range index, type mismatch) are silent no-ops;
- top-level shape is preserved: an object or array input yields a (possibly empty) object or array, while a scalar, Variant null, or the root path `$` is returned unchanged
### Why are the changes needed?
Missing feature.
### Does this PR introduce _any_ user-facing change?
Yes, new SQL function.
### How was this patch tested?
Unit tests.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code with Claude Opus 4.8
Closes #58281 from bojana-db/variant-pick.
Authored-by: bojana-db <bo...@databricks.com>
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
(cherry picked from commit ba4441f30e8c0f7718c15c7e2c208914c94f61b7)
Signed-off-by: Uros Bojanic <22...@users.noreply.github.com>
---
.../apache/spark/types/variant/VariantBuilder.java | 147 ++++++++++++++++
.../source/reference/pyspark.sql/functions.rst | 1 +
python/pyspark/sql/connect/functions/builtin.py | 13 ++
python/pyspark/sql/functions/__init__.py | 1 +
python/pyspark/sql/functions/builtin.py | 67 +++++++
python/pyspark/sql/tests/test_functions.py | 4 +
.../scala/org/apache/spark/sql/functions.scala | 47 +++++
.../sql/catalyst/analysis/FunctionRegistry.scala | 3 +-
.../variant/VariantExpressionEvalUtils.scala | 17 ++
.../expressions/variant/variantExpressions.scala | 186 ++++++++++++++++++++
.../variant/VariantExpressionSuite.scala | 194 +++++++++++++++++++++
.../apache/spark/sql/PlanGenerationTestSuite.scala | 4 +
.../explain-results/function_variant_pick.explain | 2 +
.../query-tests/queries/function_variant_pick.json | 125 +++++++++++++
.../queries/function_variant_pick.proto.bin | Bin 0 -> 1145 bytes
.../sql-functions/sql-expression-schema.md | 1 +
.../scala/org/apache/spark/sql/VariantSuite.scala | 91 ++++++++++
17 files changed, 902 insertions(+), 1 deletion(-)
diff --git a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
index ef00c8379b89..788bd02dc8b6 100644
--- a/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
+++ b/common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java
@@ -182,6 +182,29 @@ public class VariantBuilder {
return builder.result();
}
+ // Return a new variant keeping only the substructures at `paths`. A kept object field stays
+ // nested under its parent objects; kept array elements are compacted into a new array in their
+ // original order. Missing keys, out-of-range indices, and type mismatches are skipped, dropping
+ // any parent left with nothing. The top-level object/array shape is preserved (possibly empty);
+ // a scalar, variant null, or root `$` returns v unchanged.
+ public static Variant pickAtPaths(Variant v, List<PathSegment[]> paths) {
+ return pickAtPaths(v, buildPickTree(paths));
+ }
+
+ public static PickNode buildPickTree(List<PathSegment[]> paths) {
+ PickNode root = new PickNode();
+ for (PathSegment[] path : paths) {
+ root.add(path, 0);
+ }
+ return root;
+ }
+
+ public static Variant pickAtPaths(Variant v, PickNode tree) {
+ VariantBuilder builder = new VariantBuilder(false);
+ builder.pickImplTopLevel(v.value, v.metadata, v.pos, tree);
+ return builder.result();
+ }
+
// Build the variant metadata from `dictionaryKeys` and return the variant result.
public Variant result() {
int numKeys = dictionaryKeys.size();
@@ -958,6 +981,130 @@ public class VariantBuilder {
}
}
+ // A node in the keep-tree built from `variant_pick`'s JSONPaths. It records whether a path
+ // terminates here (keep everything below), else which object keys / array indices to descend
+ // into. Holds only strings and ints, so it is `Serializable` for whole-stage codegen.
+ public static final class PickNode implements java.io.Serializable {
+ // A path terminates here: keep the whole value, subsuming any deeper paths under this node.
+ private boolean keepAll = false;
+ // Lazily created; a node may hold both maps when paths disagree on the container type, and
+ // only the map matching the actual value is used.
+ private HashMap<String, PickNode> objectChildren = null;
+ private HashMap<Integer, PickNode> arrayChildren = null;
+
+ // Insert the path suffix `path[depth..]` under this node.
+ private void add(PathSegment[] path, int depth) {
+ // A broader path already keeps everything here, so any deeper path is subsumed.
+ if (keepAll) {
+ return;
+ }
+ if (depth == path.length) {
+ keepAll = true;
+ // Drop children from narrower paths added earlier; keepAll subsumes them.
+ objectChildren = null;
+ arrayChildren = null;
+ return;
+ }
+ PathSegment seg = path[depth];
+ PickNode child;
+ if (seg instanceof ObjectKeySegment) {
+ if (objectChildren == null) {
+ objectChildren = new HashMap<>();
+ }
+ child = objectChildren.computeIfAbsent(((ObjectKeySegment) seg).key, k -> new PickNode());
+ } else {
+ if (arrayChildren == null) {
+ arrayChildren = new HashMap<>();
+ }
+ child = arrayChildren.computeIfAbsent(((ArrayIndexSegment) seg).index, k -> new PickNode());
+ }
+ child.add(path, depth + 1);
+ }
+ }
+
+ // Top-level entry for `pickAtPaths`: an object or array input yields a (possibly empty) object
+ // or array; a scalar, variant null, or root `$` (`keepAll`) keeps the value unchanged.
+ private void pickImplTopLevel(byte[] value, byte[] metadata, int pos, PickNode root) {
+ checkIndex(pos, value.length);
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (root.keepAll || (basicType != OBJECT && basicType != ARRAY)) {
+ appendVariantImpl(value, metadata, pos);
+ } else if (!pickImpl(value, metadata, pos, root)) {
+ // `pickImpl` writes nothing when it keeps nothing, so the top-level empty object/array (the
+ // preserved shape) is emitted here rather than in the recursion.
+ if (basicType == OBJECT) {
+ finishWritingObject(writePos, new ArrayList<>());
+ } else {
+ finishWritingArray(writePos, new ArrayList<>());
+ }
+ }
+ }
+
+ // Append the substructures at `pos` selected by `node`, returning whether anything was appended.
+ // It writes bytes if it returns true, so misses and empty sub-containers leave no trace and the
+ // top-level empty object/array comes from `pickImplTopLevel`. A key is registered only if kept.
+ private boolean pickImpl(byte[] value, byte[] metadata, int pos, PickNode node) {
+ checkIndex(pos, value.length);
+ if (node.keepAll) {
+ appendVariantImpl(value, metadata, pos);
+ return true;
+ }
+ int basicType = value[pos] & BASIC_TYPE_MASK;
+ if (basicType == OBJECT) {
+ return handleObject(
+ value, pos, (size, idSize, offsetSize, idStart, offsetStart, dataStart) -> {
+ ArrayList<FieldEntry> fields = new ArrayList<>();
+ int start = writePos;
+ // No object-key children here: skip the whole scan and its per-field key lookups.
+ if (node.objectChildren != null) {
+ for (int i = 0; i < size; ++i) {
+ int id = readUnsigned(value, idStart + idSize * i, idSize);
+ String fieldKey = getMetadataKey(metadata, id);
+ PickNode child = node.objectChildren.get(fieldKey);
+ if (child != null) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+ int fieldOffset = writePos - start;
+ if (pickImpl(value, metadata, dataStart + offset, child)) {
+ fields.add(new FieldEntry(fieldKey, addKey(fieldKey), fieldOffset));
+ }
+ }
+ }
+ }
+ if (fields.isEmpty()) {
+ return false;
+ }
+ finishWritingObject(start, fields);
+ return true;
+ });
+ } else if (basicType == ARRAY) {
+ return handleArray(value, pos, (size, offsetSize, offsetStart, dataStart) -> {
+ ArrayList<Integer> offsets = new ArrayList<>();
+ int start = writePos;
+ // No array-index children here: skip the whole element scan.
+ if (node.arrayChildren != null) {
+ for (int i = 0; i < size; ++i) {
+ PickNode child = node.arrayChildren.get(i);
+ if (child != null) {
+ int offset = readUnsigned(value, offsetStart + offsetSize * i, offsetSize);
+ int elementOffset = writePos - start;
+ if (pickImpl(value, metadata, dataStart + offset, child)) {
+ offsets.add(elementOffset);
+ }
+ }
+ }
+ }
+ if (offsets.isEmpty()) {
+ return false;
+ }
+ finishWritingArray(start, offsets);
+ return true;
+ });
+ } else {
+ // The value is a scalar but the node still has segments to follow: nothing matches.
+ return false;
+ }
+ }
+
// Append the variant value without rewriting or creating any metadata. This is used when
// building an object during shredding, where there is a fixed pre-existing metadata that
// all shredded values will refer to.
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst b/python/docs/source/reference/pyspark.sql/functions.rst
index 94c5a9a87e13..85f6c0d51c29 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -617,6 +617,7 @@ VARIANT Functions
variant_from_arrays
variant_from_entries
variant_strip_nulls
+ variant_pick
variant_get
variant_insert
try_variant_insert
diff --git a/python/pyspark/sql/connect/functions/builtin.py b/python/pyspark/sql/connect/functions/builtin.py
index 922c8f26e0ae..72f80846ecfe 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -2343,6 +2343,19 @@ def variant_strip_nulls(v: "ColumnOrName", include_arrays: bool = True) -> Colum
variant_strip_nulls.__doc__ = pysparkfuncs.variant_strip_nulls.__doc__
+def variant_pick(v: "ColumnOrName", *paths: Union[Column, str]) -> Column:
+ if len(paths) == 0:
+ raise PySparkValueError(
+ errorClass="CANNOT_BE_EMPTY",
+ messageParameters={"item": "paths"},
+ )
+ cols = [p if isinstance(p, Column) else lit(p) for p in paths]
+ return _invoke_function("variant_pick", _to_col(v), *cols)
+
+
+variant_pick.__doc__ = pysparkfuncs.variant_pick.__doc__
+
+
def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str) -> Column:
assert isinstance(path, (Column, str))
if isinstance(path, str):
diff --git a/python/pyspark/sql/functions/__init__.py b/python/pyspark/sql/functions/__init__.py
index 861f5902f977..b12c7f997fd4 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -496,6 +496,7 @@ __all__ = [ # noqa: F405
"variant_from_arrays",
"variant_from_entries",
"variant_strip_nulls",
+ "variant_pick",
"variant_get",
"variant_insert",
"try_variant_insert",
diff --git a/python/pyspark/sql/functions/builtin.py b/python/pyspark/sql/functions/builtin.py
index 54697ac531e3..a6679140064f 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -23947,6 +23947,73 @@ def variant_strip_nulls(v: "ColumnOrName", include_arrays: bool = True) -> Colum
)
+@_try_remote_functions
+def variant_pick(v: "ColumnOrName", *paths: Union[Column, str]) -> Column:
+ """
+ Keeps only the fields or array elements of a variant at the given JSONPath locations, preserving
+ their enclosing structure; kept array elements are compacted into a new array in their
+ original order. If no path matches, an object or array input yields an empty object or array,
+ while a scalar or variant-null input is unchanged. Returns NULL if `v` is NULL; NULL paths are
+ skipped.
+
+ .. versionadded:: 4.4.0
+
+ Parameters
+ ----------
+ v : :class:`~pyspark.sql.Column` or str
+ a variant column or column name
+ paths : :class:`~pyspark.sql.Column` or str
+ one or more JSONPaths identifying substructures to keep. A `str` is a literal path; a
+ :class:`~pyspark.sql.Column` supplies the path at runtime. A valid path should start with
+ `$` and is followed by zero or more segments like `[123]`, `.name`, `['name']`, or
+ `["name"]`.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ a variant column keeping only the specified paths
+
+ Examples
+ --------
+ >>> from pyspark.sql.functions import lit, parse_json, to_json, variant_pick
+ >>> df = spark.createDataFrame([{
+ ... 'json': '''{ "a": {"b": 1, "c": 2}, "items": [10, 20, 30, 40] }''',
+ ... 'path': '$.a.b'
+ ... }])
+ >>> v = parse_json(df.json)
+ >>> df.select(to_json(variant_pick(v, "$.a.b")).alias("r")).collect()
+ [Row(r='{"a":{"b":1}}')]
+ >>> df.select(to_json(variant_pick(v, lit(None), "$.a.c", "$.items[0]")).alias("r")).collect()
+ [Row(r='{"a":{"c":2},"items":[10]}')]
+ >>> df.select(to_json(variant_pick(v, "$.items[0]", "$.items[2]")).alias("r")).collect()
+ [Row(r='{"items":[10,30]}')]
+ >>> df.select(to_json(variant_pick(v, df.path)).alias("r")).collect()
+ [Row(r='{"a":{"b":1}}')]
+ >>> df.select(to_json(variant_pick(v, "$.missing")).alias("r")).collect()
+ [Row(r='{}')]
+ >>> df.select(to_json(variant_pick(parse_json(lit('[1, 2, 3]')), "$[9]")).alias("r")).collect()
+ [Row(r='[]')]
+ >>> df.select(variant_pick(lit(None), "$.a").alias("r")).collect()
+ [Row(r=None)]
+ """
+ from pyspark.sql.classic.column import _to_java_column, _to_seq
+
+ if len(paths) == 0:
+ raise PySparkValueError(
+ errorClass="CANNOT_BE_EMPTY",
+ messageParameters={"item": "paths"},
+ )
+ sc = _get_active_spark_context()
+
+ path_cols = [p if isinstance(p, Column) else lit(p) for p in paths]
+ return _invoke_function(
+ "variant_pick",
+ _to_java_column(v),
+ _to_java_column(path_cols[0]),
+ _to_seq(sc, path_cols[1:], _to_java_column),
+ )
+
+
@_try_remote_functions
def variant_get(v: "ColumnOrName", path: Union[Column, str], targetType: str) -> Column:
"""
diff --git a/python/pyspark/sql/tests/test_functions.py b/python/pyspark/sql/tests/test_functions.py
index 594e9d4cf870..8df5538a8250 100644
--- a/python/pyspark/sql/tests/test_functions.py
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -3670,6 +3670,10 @@ class FunctionsTestsMixin:
df.select(F.to_json(F.variant_strip_nulls(strip_v, False))),
['{"a":1,"c":[1,null]}', '{"a":1,"c":[1,null]}'],
)
+ check(df.select(F.to_json(F.variant_pick(v, "$.a"))), ['{"a":1}', "{}"])
+ check(df.select(F.to_json(F.variant_pick(v, df.path))), ['{"a":1}', '{"b":2}'])
+ check(df.select(F.to_json(F.variant_pick(v, F.lit(None)))), ["{}", "{}"])
+ check(df.select(F.to_json(F.variant_pick(v, "$.a", df.path))), ['{"a":1}', '{"b":2}'])
check(df.select(F.schema_of_variant(v)), ["OBJECT<a: BIGINT>", "OBJECT<b: BIGINT>"])
check(df.select(F.schema_of_variant_agg(v)), ["OBJECT<a: BIGINT, b: BIGINT>"])
diff --git a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
index 837f4f369845..700c61cd20eb 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/functions.scala
@@ -15321,6 +15321,53 @@ object functions {
def variant_strip_nulls(v: Column, includeArrays: Boolean): Column =
Column.fn("variant_strip_nulls", v, lit(includeArrays))
+ /**
+ * Keeps only the fields or array elements of a variant at the given JSONPath locations,
+ * preserving their enclosing structure; kept array elements are compacted into a new array in
+ * their original order. If no path matches, an object or array input yields an empty object or
+ * array, while a scalar or variant-null input is unchanged. Returns NULL if `v` is NULL; NULL
+ * paths are skipped.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the column containing the first JSONPath string identifying a substructure to keep. A valid
+ * path should start with `$` and is followed by zero or more segments like `[123]`, `.name`,
+ * `['name']`, or `["name"]`.
+ * @param paths
+ * additional JSONPath arguments.
+ * @group variant_funcs
+ * @since 4.4.0
+ * @return
+ * Returns a column that evaluates to a variant.
+ */
+ @scala.annotation.varargs
+ def variant_pick(v: Column, path: Column, paths: Column*): Column =
+ Column.fn("variant_pick", (v +: path +: paths): _*)
+
+ /**
+ * Keeps only the fields or array elements of a variant at the given JSONPath locations,
+ * preserving their enclosing structure; kept array elements are compacted into a new array in
+ * their original order. If no path matches, an object or array input yields an empty object or
+ * array, while a scalar or variant-null input is unchanged. Returns NULL if `v` is NULL; NULL
+ * paths are skipped.
+ *
+ * @param v
+ * a variant column.
+ * @param path
+ * the first JSONPath identifying a substructure to keep. A valid path should start with `$`
+ * and is followed by zero or more segments like `[123]`, `.name`, `['name']`, or `["name"]`.
+ * @param paths
+ * additional JSONPath strings.
+ * @group variant_funcs
+ * @since 4.4.0
+ * @return
+ * Returns a column that evaluates to a variant.
+ */
+ @scala.annotation.varargs
+ def variant_pick(v: Column, path: String, paths: String*): Column =
+ Column.fn("variant_pick", (v +: lit(path) +: paths.map(lit)): _*)
+
/**
* Extracts a sub-variant from `v` according to `path` string, and then cast the sub-variant to
* `targetType`. Returns null if the path does not exist. Throws an exception if the cast fails.
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
index a55f4f089a78..a31e7ac4d74c 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala
@@ -1048,7 +1048,8 @@ object FunctionRegistry {
expressionBuilder("try_variant_set", TryVariantSetExpressionBuilder),
expressionBuilder("variant_array_append", VariantArrayAppendExpressionBuilder),
expressionBuilder("try_variant_array_append", TryVariantArrayAppendExpressionBuilder),
- expressionBuilder("variant_strip_nulls", VariantStripNullsExpressionBuilder)
+ expressionBuilder("variant_strip_nulls", VariantStripNullsExpressionBuilder),
+ expression[VariantPick]("variant_pick")
)
private def spatialExpressions: Seq[FunctionRegistryEntry] = Seq(
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
index b4ba27b79860..c0ce22f3b3a7 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionEvalUtils.scala
@@ -276,6 +276,23 @@ object VariantExpressionEvalUtils {
new VariantVal(out.getValue, out.getMetadata)
}
+ def parsePickPath(path: UTF8String): Array[VariantBuilder.PathSegment] =
+ toJavaSegments(parseVariantPath(path.toString, "variant_pick", allowRoot = true))
+
+ def pickAtPaths(
+ input: VariantVal,
+ paths: java.util.List[Array[VariantBuilder.PathSegment]]): VariantVal = {
+ val v = new Variant(input.getValue, input.getMetadata)
+ val out = VariantBuilder.pickAtPaths(v, paths)
+ new VariantVal(out.getValue, out.getMetadata)
+ }
+
+ def pickAtPaths(input: VariantVal, tree: VariantBuilder.PickNode): VariantVal = {
+ val v = new Variant(input.getValue, input.getMetadata)
+ val out = VariantBuilder.pickAtPaths(v, tree)
+ new VariantVal(out.getValue, out.getMetadata)
+ }
+
/** Cast a Spark value from `dataType` into the variant type. */
def castToVariant(input: Any, dataType: DataType): VariantVal = {
// Enforce strict check because it is illegal for input struct/map/variant to contain duplicate
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 e77ad10ba487..0affceb09f1e 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
@@ -1000,6 +1000,192 @@ object VariantDelete {
}
}
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(v, path1[, path2, ...]) - Keeps only the fields or array elements of a variant " +
+ "at the given JSONPath locations, preserving their enclosing structure; kept array elements " +
+ "are compacted into a new array in their original order. If no path matches, an object or " +
+ "array input yields an empty object or array, while a scalar or variant-null input is " +
+ "unchanged. Returns NULL if `v` is NULL; NULL paths are skipped.",
+ arguments = """
+ Arguments:
+ * v - A variant value to project.
+ * path1, path2, ... - One or more string expressions, each evaluating to a JSONPath
+ identifying a substructure to keep. A valid path should start with `$` and is followed by
+ zero or more segments like `[123]`, `.name`, `['name']`, or `["name"]`.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(parse_json('{"a": 1, "b": 2, "c": 3}'), '$.a', '$.c');
+ {"a":1,"c":3}
+ > SELECT _FUNC_(parse_json('{"a": {"b": 1, "c": 2}, "d": 3}'), '$.a.b');
+ {"a":{"b":1}}
+ > SELECT _FUNC_(parse_json('[10, 20, 30, 40]'), '$[0]', '$[2]');
+ [10,30]
+ > SELECT _FUNC_(parse_json('{"a": 1, "b": 2}'), NULL, '$.a', '$.missing');
+ {"a":1}
+ > SELECT _FUNC_(parse_json('{"a": {"b": 1}}'), '$.a.x');
+ {}
+ > SELECT _FUNC_(parse_json('42'), '$.a');
+ 42
+ > SELECT _FUNC_(NULL, '$.a');
+ NULL
+ """,
+ since = "4.4.0",
+ group = "variant_funcs"
+)
+// scalastyle:on line.size.limit
+case class VariantPick(children: Seq[Expression])
+ extends Expression
+ with ExpectsInputTypes {
+
+ override def dataType: DataType = VariantType
+
+ override def nullable: Boolean = children.headOption.forall(_.nullable)
+
+ override def inputTypes: Seq[AbstractDataType] = {
+ // First argument is the variant; subsequent arguments are JSONPath strings.
+ VariantType +: Seq.fill(math.max(children.length - 1, 0))(
+ StringTypeWithCollation(supportsTrimCollation = true))
+ }
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ if (children.length < 2) {
+ throw QueryCompilationErrors.wrongNumArgsError(
+ prettyName, Seq("> 1"), children.length)
+ }
+ super.checkInputDataTypes()
+ }
+
+ private def variantChild: Expression = children.head
+ private def pathChildren: Seq[Expression] = children.tail
+
+ @transient private lazy val pathArgs: Seq[VariantPick.PickPathArg] =
+ pathChildren.flatMap(VariantPick.toPathArg)
+
+ // When every path is constant, the keep-tree is the same for all rows, so build it once here and
+ // reuse it rather than rebuilding it per row. `None` means at least one path is dynamic.
+ @transient private lazy val foldableTree: Option[VariantBuilder.PickNode] = {
+ if (pathArgs.forall(_.isInstanceOf[VariantPick.ParsedPickPath])) {
+ val paths = new java.util.ArrayList[Array[VariantBuilder.PathSegment]](pathArgs.length)
+ pathArgs.foreach {
+ case parsed: VariantPick.ParsedPickPath => paths.add(parsed.javaSegments)
+ case _ =>
+ }
+ Some(VariantBuilder.buildPickTree(paths))
+ } else {
+ None
+ }
+ }
+
+ // Collect the java path segments of all non-NULL paths.
+ private def collectPaths(
+ input: InternalRow): java.util.List[Array[VariantBuilder.PathSegment]] = {
+ val paths = new java.util.ArrayList[Array[VariantBuilder.PathSegment]](pathArgs.length)
+ var i = 0
+ while (i < pathArgs.length) {
+ pathArgs(i) match {
+ case parsed: VariantPick.ParsedPickPath =>
+ paths.add(parsed.javaSegments)
+ case VariantPick.DynamicPickPath(expr) =>
+ val pathVal = expr.eval(input).asInstanceOf[UTF8String]
+ if (pathVal != null) {
+ paths.add(VariantExpressionEvalUtils.parsePickPath(pathVal))
+ }
+ }
+ i += 1
+ }
+ paths
+ }
+
+ override def eval(input: InternalRow): Any = {
+ // Force constant-path parsing before the NULL short-circuit so a malformed constant path fails
+ // identically in interpreted and codegen modes, even when the input variant is NULL (codegen
+ // parses constant paths at code-generation time).
+ val tree = foldableTree
+ val inputVariant = variantChild.eval(input).asInstanceOf[VariantVal]
+ if (inputVariant == null) return null
+ tree match {
+ case Some(t) => VariantExpressionEvalUtils.pickAtPaths(inputVariant, t)
+ case None => VariantExpressionEvalUtils.pickAtPaths(inputVariant, collectPaths(input))
+ }
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
+ val cls = VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
+ val variantValType = CodeGenerator.javaType(VariantType)
+ val childCode = variantChild.genCode(ctx)
+
+ val project = foldableTree match {
+ case Some(tree) =>
+ // Constant paths: use the keep-tree built once and shipped as a reference.
+ val treeArg = ctx.addReferenceObj("vpTree", tree, classOf[VariantBuilder.PickNode].getName)
+ s"${ev.value} = $cls.pickAtPaths(${childCode.value}, $treeArg);"
+ case None =>
+ // At least one dynamic path: gather the paths per row, then project.
+ val paths = ctx.freshName("vpPaths")
+ val addPaths = pathArgs.map {
+ case parsed: VariantPick.ParsedPickPath =>
+ val parsedArg = ctx.addReferenceObj("vpParsed", parsed)
+ s"$paths.add($parsedArg.javaSegments());"
+ case VariantPick.DynamicPickPath(expr) =>
+ val pCode = expr.genCode(ctx)
+ s"""
+ |${pCode.code}
+ |if (!${pCode.isNull}) {
+ | $paths.add($cls.parsePickPath(${pCode.value}));
+ |}
+ """.stripMargin
+ }.mkString("\n")
+ s"""
+ |java.util.ArrayList $paths = new java.util.ArrayList();
+ |$addPaths
+ |${ev.value} = $cls.pickAtPaths(${childCode.value}, $paths);
+ """.stripMargin
+ }
+
+ val code = code"""
+ ${childCode.code}
+ boolean ${ev.isNull} = ${childCode.isNull};
+ $variantValType ${ev.value} = ${CodeGenerator.defaultValue(VariantType)};
+ if (!${ev.isNull}) {
+ $project
+ }
+ """
+ ev.copy(code = code)
+ }
+
+ override def prettyName: String = "variant_pick"
+
+ override protected def withNewChildrenInternal(
+ newChildren: IndexedSeq[Expression]): VariantPick = copy(children = newChildren)
+}
+
+object VariantPick {
+ sealed trait PickPathArg
+ case class ParsedPickPath(segments: Array[VariantPathSegment]) extends PickPathArg {
+ // `VariantBuilder.PathSegment` is not `Serializable`, so the cached Java form is
+ // `@transient` and re-initialized once per executor task after deserialization.
+ @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
+ VariantExpressionEvalUtils.toJavaSegments(segments)
+ }
+ case class DynamicPickPath(expr: Expression) extends PickPathArg
+
+ private[variant] def toPathArg(child: Expression): Option[PickPathArg] = {
+ if (child.foldable) {
+ val v = child.eval()
+ if (v == null) {
+ None
+ } else {
+ Some(ParsedPickPath(VariantExpressionEvalUtils.parseVariantPath(
+ v.asInstanceOf[UTF8String].toString, "variant_pick", allowRoot = true)))
+ }
+ } else {
+ Some(DynamicPickPath(child))
+ }
+ }
+}
+
case class VariantInsert(
input: Expression,
path: Expression,
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
index ed790e643adc..75018a572b8b 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala
@@ -2127,4 +2127,198 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
StringType),
null)
}
+
+ test("variant_pick") {
+ def checkPick(input: String, paths: Seq[String], expected: String): Unit = {
+ val pathLits: Seq[Expression] = paths.map(p => Literal.create(p, StringType))
+ val expr = VariantPick(Literal(parseJson(input)) +: pathLits)
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(expr, StringType)),
+ expected)
+ }
+
+ // Keep a subset of object fields, and union paths under the same parent (parent preserved).
+ checkPick("""{"a": 1, "b": 2, "c": 3}""", Seq("$.a", "$.c"), """{"a":1,"c":3}""")
+ checkPick("""{"a": {"b": 1, "c": 2}, "d": 3}""", Seq("$.a.b"), """{"a":{"b":1}}""")
+ checkPick(
+ """{"a": {"b": 1, "c": 2, "d": 3}}""", Seq("$.a.b", "$.a.c"), """{"a":{"b":1,"c":2}}""")
+
+ // A broader path subsumes a narrower one.
+ checkPick("""{"a": {"b": 1, "c": 2}}""", Seq("$.a", "$.a.b"), """{"a":{"b":1,"c":2}}""")
+ checkPick("""{"a": {"b": 1, "c": 2}}""", Seq("$.a.b", "$.a"), """{"a":{"b":1,"c":2}}""")
+
+ // Arrays are compacted in original order; out-of-range indices are skipped.
+ checkPick("[10, 20, 30, 40]", Seq("$[2]", "$[0]"), "[10,30]")
+ checkPick("[10, 20, 30]", Seq("$[5]"), "[]")
+ checkPick("""{"a": [10, 20, 30]}""", Seq("$.a[1]"), """{"a":[20]}""")
+
+ // Array elements that are containers recurse and compact; an element that picks nothing drops.
+ checkPick("""[{"x": 1, "y": 2}, {"z": 3}]""", Seq("$[0].x"), """[{"x":1}]""")
+ checkPick("""[{"x": 1}]""", Seq("$[0].y"), "[]")
+ checkPick("[[10, 20, 30]]", Seq("$[0][1]"), "[[20]]")
+
+ // Missing keys are skipped; a parent whose picks all miss is dropped.
+ checkPick("""{"a": 1, "b": 2}""", Seq("$.a", "$.missing"), """{"a":1}""")
+ checkPick("""{"a": {"b": 1}}""", Seq("$.a.x"), "{}")
+
+ // Root path is identity; a scalar or variant-null input is returned unchanged. The variant
+ // null is compared directly, since casting it to a string yields SQL NULL, not the text "null".
+ checkPick("""{"a": 1}""", Seq("$"), """{"a":1}""")
+ checkPick("42", Seq("$.a"), "42")
+ checkEvaluation(
+ VariantPick(Seq(Literal(parseJson("null")), Literal("$.a"))),
+ parseJson("null"))
+
+ // Type mismatch matches nothing, but the top-level shape is preserved (object -> {}, array
+ // -> []); descending into a scalar, or into a container of the wrong kind, drops the parent.
+ checkPick("[1, 2, 3]", Seq("$.a"), "[]")
+ checkPick("""{"a": 1}""", Seq("$[0]"), "{}")
+ checkPick("""{"a": 1}""", Seq("$.a.b"), "{}")
+ checkPick("""{"a": [1, 2]}""", Seq("$.a.b"), "{}")
+ checkPick("""{"a": {"b": 1}}""", Seq("$.a[0]"), "{}")
+
+ // Both-maps node uses only the matching branch.
+ checkPick("""{"a": 1, "b": 2}""", Seq("$.a", "$[0]"), """{"a":1}""")
+ checkPick("[10, 20, 30]", Seq("$.a", "$[0]"), "[10]")
+ checkPick("""{"a": {"b": 1}}""", Seq("$.a.b", "$.a[0]", "$.a"), """{"a":{"b":1}}""")
+
+ // Duplicate paths are deduplicated, and bracket notation resolves like dot notation.
+ checkPick("""{"a": 1, "b": 2}""", Seq("$.a", "$.a"), """{"a":1}""")
+ checkPick("""{"a": 1, "b": 2}""", Seq("$['a']"), """{"a":1}""")
+
+ // Non-ASCII field names are matched and preserved, in both dot and bracket notation.
+ // scalastyle:off nonascii
+ checkPick("""{"café": 1, "naïve": 2}""", Seq("$.café"), """{"café":1}""")
+ checkPick(
+ """{"日本語": {"x": 1, "y": 2}}""", Seq("$['日本語'].x"), """{"日本語":{"x":1}}""")
+ checkPick("""{"ключ": [10, 20, 30]}""", Seq("$.ключ[1]"), """{"ключ":[20]}""")
+ // scalastyle:on nonascii
+
+ // Larger object trees built from several paths.
+ // Many sibling keys unioned under one parent.
+ checkPick(
+ """{"a": {"b": 1, "c": 2, "d": 3, "e": 4, "f": 5, "g": 6, "h": 7}}""",
+ Seq("$.a.b", "$.a.c", "$.a.e", "$.a.g", "$.a.h"),
+ """{"a":{"b":1,"c":2,"e":4,"g":6,"h":7}}""")
+ // Several top-level parents, each keeping a different subset of children.
+ checkPick(
+ """{"a": {"x": 1, "y": 2}, "b": {"p": 3, "q": 4}, "c": 5, "d": {"m": 6, "n": 7}}""",
+ Seq("$.a.x", "$.b.p", "$.b.q", "$.c", "$.d.m"),
+ """{"a":{"x":1},"b":{"p":3,"q":4},"c":5,"d":{"m":6}}""")
+ // A broader path subsumes a narrower sibling while other branches are kept partially.
+ checkPick(
+ """{"a": {"b": {"c": 1, "d": 2}, "e": 3}, "f": {"g": 4, "h": 5}}""",
+ Seq("$.a.b.c", "$.a.b", "$.a.e", "$.f.g"),
+ """{"a":{"b":{"c":1,"d":2},"e":3},"f":{"g":4}}""")
+ // Deep branch: multiple paths share a long prefix; an unpicked sibling is dropped.
+ checkPick(
+ """{"a": {"b": {"c": {"d": 1, "e": 2}, "f": 3, "g": 4}}}""",
+ Seq("$.a.b.c.d", "$.a.b.c.e", "$.a.b.f"),
+ """{"a":{"b":{"c":{"d":1,"e":2},"f":3}}}""")
+ // A path terminating at a genuinely empty container keeps it as-is (not dropped).
+ checkPick("""{"a": {}, "b": 1}""", Seq("$.a"), """{"a":{}}""")
+ // A picked field that matches nothing is rewound.
+ checkPick(
+ """{"a": 1, "b": {"m": 1}, "c": 2}""",
+ Seq("$.a", "$.b.x", "$.c"),
+ """{"a":1,"c":2}""")
+
+ // Larger array trees.
+ // Indices given out of order are compacted into original array order, not path order.
+ checkPick(
+ "[0, 10, 20, 30, 40, 50, 60, 70]",
+ Seq("$[5]", "$[0]", "$[7]", "$[3]"),
+ "[0,30,50,70]")
+ // A picked element that matches nothing is dropped; survivors compact around the hole.
+ checkPick(
+ """[{"x": 1}, {"y": 2}, {"z": 3}]""",
+ Seq("$[0].x", "$[1].w", "$[2].z"),
+ """[{"x":1},{"z":3}]""")
+ // Nested arrays: picks descend into multiple sub-arrays, compacting at each level.
+ checkPick(
+ "[[10, 20], [30, 40], [50, 60]]",
+ Seq("$[0][1]", "$[2][0]"),
+ "[[20],[50]]")
+ // A path ending at an array index keeps the whole element; unpicked siblings drop.
+ checkPick(
+ """[{"x": 1, "y": 2}, {"z": 3}]""",
+ Seq("$[0]"),
+ """[{"x":1,"y":2}]""")
+
+ // Mixed object/array variant with paths at different depths.
+ checkPick(
+ """{"a": [{"p": 1, "q": 2, "r": 3}, {"p": 4, "q": 5, "r": 6}], "b": {"s": 7, "t": 8}}""",
+ Seq("$.a[0].p", "$.a[0].q", "$.a[1].r", "$.b.s"),
+ """{"a":[{"p":1,"q":2},{"r":6}],"b":{"s":7}}""")
+
+ // A large variant with many paths at once.
+ checkPick(
+ """{"a": {"b": 1, "c": 2, "d": 3},
+ "e": [{"f": 10, "g": 20}, {"f": 30, "g": 40}, {"f": 50}],
+ "h": {"i": {"j": 5, "k": 6}}, "l": 7}""",
+ Seq("$.a.b", "$.a.c", "$.e[0].f", "$.e[2]", "$.h.i.j", "$.h.i", "$.l", "$.x"),
+ """{"a":{"b":1,"c":2},"e":[{"f":10},{"f":50}],"h":{"i":{"j":5,"k":6}},"l":7}""")
+
+ // Deep, branchy trie exercising many branches at once.
+ checkPick(
+ """{"a": {"b": [{"c": 1, "d": [10, 11, 12], "e": {"f": 1, "g": 2}}, {"c": 2},
+ {"e": {"f": 5}}], "h": {"i": {"j": 9}, "l": 8}}}""",
+ Seq("$.a.b[0].c", "$.a.b[0].d[1]", "$.a.b[0].e.f", "$.a.b[2].e", "$.a.b[1].z", "$.a.b.x",
+ "$.a.h.i", "$.a.h.i.j"),
+ """{"a":{"b":[{"c":1,"d":[11],"e":{"f":1}},{"e":{"f":5}}],"h":{"i":{"j":9}}}}""")
+
+ // A NULL path is skipped; the remaining paths still apply.
+ checkPick("""{"a": 1, "b": 2}""", Seq(null, "$.a"), """{"a":1}""")
+
+ // Dynamic (non-foldable) path mixed with a literal.
+ val mixedLitDyn = VariantPick(Seq(
+ Literal(parseJson("""{"a": 1, "b": 2, "c": 3}""")),
+ Literal("$.a"),
+ BoundReference(0, StringType, nullable = true)))
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(mixedLitDyn, StringType)),
+ """{"a":1,"c":3}""",
+ InternalRow(UTF8String.fromString("$.c")))
+
+ // A dynamic path that evaluates to NULL is skipped at runtime (unlike a constant NULL, which is
+ // dropped when the tree is built); the remaining literal path still applies.
+ checkEvaluation(
+ ResolveTimeZone.resolveTimeZones(Cast(mixedLitDyn, StringType)),
+ """{"a":1}""",
+ InternalRow(null))
+
+ // NULL variant input yields NULL.
+ checkEvaluation(
+ VariantPick(Seq(Literal.create(null, VariantType), Literal("$.a"))),
+ null)
+
+ // Malformed paths are rejected.
+ checkErrorInExpression[SparkRuntimeException](
+ VariantPick(Seq(Literal(parseJson("""{"a": 1}""")), Literal("garbage"))),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "garbage", "functionName" -> "`variant_pick`"))
+
+ // A malformed constant path is rejected in both interpreted and codegen modes even when the
+ // input variant is NULL: the path is parsed before the NULL short-circuit.
+ checkErrorInExpression[SparkRuntimeException](
+ VariantPick(Seq(BoundReference(0, VariantType, nullable = true), Literal("garbage"))),
+ InternalRow(null),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "garbage", "functionName" -> "`variant_pick`"))
+
+ // A malformed dynamic path is rejected at runtime (via `parsePickPath`), in both modes.
+ checkErrorInExpression[SparkRuntimeException](
+ VariantPick(Seq(
+ Literal(parseJson("""{"a": 1}""")),
+ BoundReference(0, StringType, nullable = true))),
+ InternalRow(UTF8String.fromString("garbage")),
+ "INVALID_VARIANT_PATH",
+ Map("path" -> "garbage", "functionName" -> "`variant_pick`"))
+
+ // At least one path is required.
+ val noPaths = VariantPick(Seq(Literal(parseJson("""{"a": 1}"""))))
+ intercept[org.apache.spark.sql.AnalysisException] {
+ noPaths.checkInputDataTypes()
+ }
+ }
}
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
index b117f1a0b5cd..1e8267f1b0bc 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
@@ -2845,6 +2845,10 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging {
fn.variant_strip_nulls(fn.parse_json(fn.col("g")), false)
}
+ functionTest("variant_pick") {
+ fn.variant_pick(fn.parse_json(fn.col("g")), "$.a", "$.b")
+ }
+
functionTest("variant_get") {
fn.variant_get(fn.parse_json(fn.col("g")), "$", "int")
}
diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_pick.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_pick.explain
new file mode 100644
index 000000000000..6193b2fb67d2
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_pick.explain
@@ -0,0 +1,2 @@
+Project [variant_pick(static_invoke(VariantExpressionEvalUtils.parseJson(g#0, false, true, true)), $.a, $.b) AS variant_pick(parse_json(g), $.a, $.b)#0]
++- LocalRelation <empty>, [id#0L, a#0, b#0, d#0, e#0, f#0, g#0]
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.json b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.json
new file mode 100644
index 000000000000..d31ab5d81836
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.json
@@ -0,0 +1,125 @@
+{
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "localRelation": {
+ "schema": "struct\u003cid:bigint,a:int,b:double,d:struct\u003cid:bigint,a:int,b:double\u003e,e:array\u003cint\u003e,f:map\u003cstring,struct\u003cid:bigint,a:int,b:double\u003e\u003e,g:string\u003e"
+ }
+ },
+ "expressions": [{
+ "unresolvedFunction": {
+ "functionName": "variant_pick",
+ "arguments": [{
+ "unresolvedFunction": {
+ "functionName": "parse_json",
+ "arguments": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "g"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "col",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "parse_json",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string": "$.a"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "variant_pick",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "literal": {
+ "string": "$.b"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "variant_pick",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "isInternal": false
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "variant_pick",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.proto.bin
new file mode 100644
index 000000000000..f8b643b6ce78
Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/function_variant_pick.proto.bin differ
diff --git a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
index 6f11ca0435ae..2bfa746fbc6e 100644
--- a/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
+++ b/sql/core/src/test/resources/sql-functions/sql-expression-schema.md
@@ -580,6 +580,7 @@
| org.apache.spark.sql.catalyst.expressions.variant.VariantFromEntries | variant_from_entries | SELECT variant_from_entries(array(struct('a', 1), struct('b', 2))) | struct<variant_from_entries(array(struct(a, 1), struct(b, 2))):variant> |
| org.apache.spark.sql.catalyst.expressions.variant.VariantGetExpressionBuilder | variant_get | SELECT variant_get(parse_json('{"a": 1}'), '$.a', 'int') | struct<variant_get(parse_json({"a": 1}), $.a):int> |
| org.apache.spark.sql.catalyst.expressions.variant.VariantInsertExpressionBuilder | variant_insert | SELECT variant_insert(parse_json('{"a": 1}'), '$.b', 2) | struct<variant_insert(parse_json({"a": 1}), $.b, 2):variant> |
+| org.apache.spark.sql.catalyst.expressions.variant.VariantPick | variant_pick | SELECT variant_pick(parse_json('{"a": 1, "b": 2, "c": 3}'), '$.a', '$.c') | struct<variant_pick(parse_json({"a": 1, "b": 2, "c": 3}), $.a, $.c):variant> |
| org.apache.spark.sql.catalyst.expressions.variant.VariantSetExpressionBuilder | variant_set | SELECT variant_set(parse_json('{"a": 1}'), '$.a', 2) | struct<variant_set(parse_json({"a": 1}), $.a, 2, true):variant> |
| org.apache.spark.sql.catalyst.expressions.variant.VariantStripNullsExpressionBuilder | variant_strip_nulls | SELECT variant_strip_nulls(parse_json('{"a": 1, "b": null, "c": 3}')) | struct<variant_strip_nulls(parse_json({"a": 1, "b": null, "c": 3}), true):variant> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathBoolean | xpath_boolean | SELECT xpath_boolean('<a><b>1</b></a>','a/b') | struct<xpath_boolean(<a><b>1</b></a>, a/b):boolean> |
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 9c2bfaca03a2..765edf1de809 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
@@ -761,6 +761,97 @@ class VariantSuite extends SharedSparkSession with ExpressionEvalHelper {
rows("""{"a":[1,null]}""", """{"y":2}""", null))
}
+ test("variant_pick with literal paths") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+
+ // SQL parse/registration path end to end, including array compaction and the root identity.
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(parse_json('{\"a\": 1, \"b\": 2, \"c\": 3}'), " +
+ "'$.a', '$.c'))"),
+ rows("""{"a":1,"c":3}"""))
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(parse_json('[10, 20, 30, 40]'), '$[2]', '$[0]'))"),
+ rows("[10,30]"))
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(parse_json('{\"a\": 1}'), '$'))"),
+ rows("""{"a":1}"""))
+
+ // The String* DataFrame-API overload compiles each path to a literal.
+ val df = Seq("""{"a": 1, "b": 2, "c": 3}""").toDF("json")
+ checkAnswer(
+ df.select(to_json(variant_pick(parse_json(col("json")), "$.a", "$.c")).alias("r")),
+ rows("""{"a":1,"c":3}"""))
+
+ // A NULL path is skipped; a fully-missing projection yields an empty container.
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(parse_json('{\"a\": 1, \"b\": 2}'), NULL, '$.a'))"),
+ rows("""{"a":1}"""))
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(parse_json('{\"a\": 1}'), '$.missing'))"),
+ rows("{}"))
+ checkAnswer(
+ sql("SELECT to_json(variant_pick(CAST(NULL AS VARIANT), '$.a'))"),
+ rows(null))
+
+ // A malformed path is rejected.
+ checkError(
+ exception = intercept[SparkRuntimeException] {
+ sql("SELECT variant_pick(parse_json('{\"a\": 1}'), 'garbage')").collect()
+ },
+ condition = "INVALID_VARIANT_PATH",
+ parameters = Map("path" -> "garbage", "functionName" -> toSQLId("variant_pick")))
+
+ // At least one path argument is required.
+ checkError(
+ exception = intercept[AnalysisException] {
+ sql("SELECT variant_pick(parse_json('{}'))").collect()
+ },
+ condition = "WRONG_NUM_ARGS.WITHOUT_SUGGESTION",
+ parameters = Map(
+ "functionName" -> toSQLId("variant_pick"),
+ "expectedNum" -> "> 1",
+ "actualNum" -> "1",
+ "docroot" -> "https://spark.apache.org/docs/latest"))
+ }
+
+ test("variant_pick with dynamic paths") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+ Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { codegenMode =>
+ withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) {
+ val df = Seq(
+ ("""{"a": 1, "b": 2, "c": 3}""", "$.c", "$.a"),
+ ("""{"a": 1, "b": 2}""", "$.b", "$.b"),
+ (null, "$.a", "$.a")
+ ).toDF("json", "p1", "p2")
+ val v = parse_json(col("json"))
+ // Dynamic (non-foldable) column paths, collected per row; row 2 has a duplicate path.
+ val out = df.select(to_json(variant_pick(v, col("p1"), col("p2"))).alias("r"))
+ checkAnswer(out, rows("""{"a":1,"c":3}""", """{"b":2}""", null))
+ }
+ }
+ }
+
+ test("variant_pick with a mix of literal and non-literal paths") {
+ def rows(results: Any*): Seq[Row] = results.map(Row(_))
+ Seq("CODEGEN_ONLY", "NO_CODEGEN").foreach { codegenMode =>
+ withSQLConf(SQLConf.CODEGEN_FACTORY_MODE.key -> codegenMode) {
+ val df = Seq(
+ ("""{"a": 1, "b": 2, "c": 3}""", "$.b", "$.c"),
+ ("""{"a": 1, "b": 2, "c": 3}""", "$.b", null)
+ ).toDF("json", "p1", "p2")
+ val v = parse_json(col("json"))
+ // One literal path followed by two dynamic column paths; the second is NULL in row 2.
+ checkAnswer(
+ df.select(to_json(variant_pick(v, lit("$.a"), col("p1"), col("p2"))).alias("r")),
+ rows("""{"a":1,"b":2,"c":3}""", """{"a":1,"b":2}"""))
+ // A dynamic column path followed by two literal paths; the dynamic is NULL in row 2.
+ checkAnswer(
+ df.select(to_json(variant_pick(v, col("p2"), lit("$.a"), lit("$.b"))).alias("r")),
+ rows("""{"a":1,"b":2,"c":3}""", """{"a":1,"b":2}"""))
+ }
+ }
+ }
+
test("round trip tests") {
withSQLConf(SQLConf.VARIANT_INFER_SHREDDING_SCHEMA.key -> "false") {
val rand = new Random(42)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org