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