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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
