You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "bojana-db (via GitHub)" <gi...@apache.org> on 2026/08/25 11:41:21 UTC
[PR] [SPARK-XXXXX][SQL] Add `variant_pick` expression [spark]
bojana-db opened a new pull request, #58281:
URL: https://github.com/apache/spark/pull/58281
<!--
Thanks for sending a pull request! Here are some tips for you:
1. If this is your first time, please read our contributor guidelines: https://spark.apache.org/contributing.html
2. Ensure you have added or run the appropriate tests for your PR: https://spark.apache.org/developer-tools.html
3. If the PR is unfinished, add '[WIP]' in your PR title, e.g., '[WIP][SPARK-XXXX] Your PR title ...'.
4. Be sure to keep the PR description updated to reflect all changes.
5. Please write your PR title to summarize what this PR proposes.
6. If possible, provide a concise example to reproduce the issue for a faster review.
7. If you want to add a new configuration, please read the guideline first for naming configurations in
'common/utils/src/main/scala/org/apache/spark/internal/config/ConfigEntry.scala'.
8. If you want to add or modify an error type or message, please read the guideline first in
'common/utils/src/main/resources/error/README.md'.
-->
### What changes were proposed in this pull request?
<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If you fix some SQL features, you can provide some references of other DBMSes.
3. If there is design documentation, please add the link.
4. If there is a discussion in the mailing list, please add the link.
-->
### Why are the changes needed?
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such as new features, bug fixes, or other behavior changes. Documentation-only updates are not considered user-facing changes.
If yes, please clarify the previous behavior and the change this PR proposes - provide the console output, description and/or an example to show the behavior difference if possible.
If possible, please also clarify if this is a user-facing change compared to the released Spark versions or within the unreleased branches such as master.
If no, write 'No'.
-->
### How was this patch tested?
<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
If benchmark tests were added, please run the benchmarks in GitHub Actions for the consistent environment, and the instructions could accord to: https://spark.apache.org/developer-tools.html#github-workflow-benchmarks.
-->
### Was this patch authored or co-authored using generative AI tooling?
<!--
If generative AI tooling has been used in the process of authoring this patch, please include the
phrase: 'Generated-by: ' followed by the name of the tool and its version.
If no, write 'No'.
Please refer to the [ASF Generative Tooling Guidance](https://www.apache.org/legal/generative-tooling.html) for details.
-->
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on PR #58281:
URL: https://github.com/apache/spark/pull/58281#issuecomment-5603220115
**Merge Summary:**
- merged into master https://github.com/apache/spark/commit/ba4441f30e8c0f7718c15c7e2c208914c94f61b7
- merged into branch-4.x https://github.com/apache/spark/commit/d8dd0bec59c0b45e4e61f34c40b317ba697549bd
*Posted by `merge_spark_pr.py`*
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b closed pull request #58281: [SPARK-59130][SQL] Add `variant_pick` expression
URL: https://github.com/apache/spark/pull/58281
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "ZdravDim (via GitHub)" <gi...@apache.org>.
ZdravDim commented on PR #58281:
URL: https://github.com/apache/spark/pull/58281#issuecomment-5531100217
.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "harshmotw-db (via GitHub)" <gi...@apache.org>.
harshmotw-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3938300418
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -23544,6 +23544,71 @@ 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='{}')]
Review Comment:
I see that we have a test for it. A documented example could help
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3950004056
##########
sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala:
##########
Review Comment:
Already added in the dynamic test group, but moved now to separate test group.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3940726815
##########
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)
Review Comment:
This is done in order to match other engines. NULL paths in variadic functions are always skipped, rather than nulling the whole output.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3950014080
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala:
##########
@@ -2127,4 +2127,123 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
StringType),
null)
}
+
+ test("variant_pick") {
Review Comment:
Added more tests.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3968347110
##########
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"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
Review Comment:
This file is generated by the `PlanGenerationTestSuite` and every such file ends with no new line. Adding it here would be inconsistent. The proper fix if we want newlines in the end of the files is to change the generator.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3949983470
##########
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)
Review Comment:
Hmm, I'm not sure how much this would actually buy us. The literal paths are already parsed once and cached. The dynamic paths have to be re-parsed every row regardless, since their value changes.
So a reusable mutable array wouldn't skip any parsing; it would only save the per-row `ArrayList` allocation, while the dominant cost (rebuilding the `PickNode` tree each row) stays.
I'd lean toward leaving it as is, but happy to discuss it further.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3960870174
##########
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)
Review Comment:
Oh, I see what you meant - it only takes first argument (`headOption`), so paths do not affect nullability.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3950008362
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -23544,6 +23544,71 @@ 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='{}')]
Review Comment:
Added.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "harshmotw-db (via GitHub)" <gi...@apache.org>.
harshmotw-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3960675480
##########
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)
Review Comment:
Yeah, I agree that null paths are skipped. That's why the nullability of the whole expression should not depend on the nullability of the paths, right? The nullability of the expression only depends on the nullability of the variant child.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "harshmotw-db (via GitHub)" <gi...@apache.org>.
harshmotw-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3938306957
##########
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)
Review Comment:
If you agree with this, we could reconsider this for VariantDelete and other expressions as well? https://github.com/bojana-db/spark/blob/310f00b494427fba6cd64d17d6473e3fadeb6e32/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala#L895
##########
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)
Review Comment:
If you agree with this, we could consider this for VariantDelete and other expressions as well? https://github.com/bojana-db/spark/blob/310f00b494427fba6cd64d17d6473e3fadeb6e32/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala#L895
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "harshmotw-db (via GitHub)" <gi...@apache.org>.
harshmotw-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3960708219
##########
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)
Review Comment:
Let's leave as is then. Was just making sure we considered it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3961033395
##########
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"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
Review Comment:
```suggestion
}
```
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "uros-b (via GitHub)" <gi...@apache.org>.
uros-b commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3961035912
##########
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"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
Review Comment:
Please end all files with a proper newline.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "harshmotw-db (via GitHub)" <gi...@apache.org>.
harshmotw-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3938074437
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -958,6 +981,125 @@ private void appendWithNullStrippingImpl(
}
}
+ // 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 {
+ pickImpl(value, metadata, pos, root);
+ }
+ }
+
+ // Append the substructures of the value at `pos` selected by `node`, and return whether anything
+ // was appended. A caller drops a field or element whose pick produced nothing by resetting
+ // `writePos`, so unmatched paths (missing keys, out-of-range indices, or type mismatches) leave
+ // no trace. A dictionary key is registered only for a field that is actually 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 fieldStart = writePos;
+ int fieldOffset = writePos - start;
+ if (pickImpl(value, metadata, dataStart + offset, child)) {
+ fields.add(new FieldEntry(fieldKey, addKey(fieldKey), fieldOffset));
+ } else {
+ writePos = fieldStart;
Review Comment:
Is this necessary? If `pickImpl` returns false (nothing was picked in the subtree), could `writePos` have been modified? Same with the array branch
**Edit:** Oh, I guess `finishWritingObject` would write empty objects in subtrees and we want to revert all of that unnecessary progress? Is there a better way to do this? Maybe only run `finishWritingObject` when fields.empty of `depth == 0` or something? Just brainstorming
##########
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)
Review Comment:
Why don't we make this a mutable array as a class variable and just modify the dynamic indices during every row?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/variant/VariantExpressionSuite.scala:
##########
@@ -2127,4 +2127,123 @@ class VariantExpressionSuite extends SparkFunSuite with ExpressionEvalHelper {
StringType),
null)
}
+
+ test("variant_pick") {
Review Comment:
I don't think we're stressing the tree creation code enough in these tests because it seems we have at most two paths in each test? If that's the case, the trees would be very small and simple
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -23544,6 +23544,71 @@ 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='{}')]
Review Comment:
If you pick a missing index/key from a variant that is a top-level array, it would return `[]` right? Can you add an example for it?
##########
sql/core/src/test/scala/org/apache/spark/sql/VariantSuite.scala:
##########
Review Comment:
Let's also add tests with a combination of literal/non-literal paths in the arguments.
##########
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)
Review Comment:
Do paths have an impact on nullability? I don't believe they do
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on PR #58281:
URL: https://github.com/apache/spark/pull/58281#issuecomment-5589693547
@uros-db Please review.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-59130][SQL] Add `variant_pick` expression [spark]
Posted by "bojana-db (via GitHub)" <gi...@apache.org>.
bojana-db commented on code in PR #58281:
URL: https://github.com/apache/spark/pull/58281#discussion_r3950160361
##########
common/variant/src/main/java/org/apache/spark/types/variant/VariantBuilder.java:
##########
@@ -958,6 +981,125 @@ private void appendWithNullStrippingImpl(
}
}
+ // 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 {
+ pickImpl(value, metadata, pos, root);
+ }
+ }
+
+ // Append the substructures of the value at `pos` selected by `node`, and return whether anything
+ // was appended. A caller drops a field or element whose pick produced nothing by resetting
+ // `writePos`, so unmatched paths (missing keys, out-of-range indices, or type mismatches) leave
+ // no trace. A dictionary key is registered only for a field that is actually 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 fieldStart = writePos;
+ int fieldOffset = writePos - start;
+ if (pickImpl(value, metadata, dataStart + offset, child)) {
+ fields.add(new FieldEntry(fieldKey, addKey(fieldKey), fieldOffset));
+ } else {
+ writePos = fieldStart;
Review Comment:
Good catch, addressed.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org