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: [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]
