bojana-db commented on code in PR #58301:
URL: https://github.com/apache/spark/pull/58301#discussion_r4073218757
##########
python/pyspark/sql/tests/test_functions.py:
##########
@@ -3595,6 +3595,9 @@ def check(resultDf, expected):
check(df.select(F.is_variant_null(v)), [False, False])
check(df.select(F.is_valid_variant(v)), [True, True])
+ check(df.select(F.variant_array_length(v)), [None, None])
+ check(df.select(F.variant_array_length(v, "$.a")), [None, None])
+ check(df.select(F.variant_array_length(v, df.path)), [None, None])
Review Comment:
Tests below are enough (testing all the paths).
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
Review Comment:
Consider using a single expression class with the path defaulting to root
(`$`) instead of two classes (`VariantArrayLength` +
`VariantArrayLengthWithPath`) dispatched by arity. `variant_array_length(v)` is
semantically `variant_array_length(v, '$')`, and `getVariant` with an empty
path returns the variant uncopied, so one class would cover both. That would
also match how `variant_strip_nulls` (a `functionSignature` default) and
`variant_get` (single class, defaulted arg) handle optional arguments.
##########
sql/connect/common/src/test/resources/query-tests/explain-results/function_variant_array_length.explain:
##########
Review Comment:
Let's test one arg variation too.
##########
sql/core/src/test/scala/org/apache/spark/sql/VariantEndToEndSuite.scala:
##########
Review Comment:
There are a few test coverage gaps we should close:
- `VariantArrayLengthWithPath` has separate interpreted (`nullSafeEval`)
and codegen (`doGenCode`) paths, but the added tests run only default
whole-stage codegen. This can be achieved through `checkEvaluation` - see
`VariantExpressionSuite.scala` - you can add there tests for
`variant_array_length` similarly to other variant expressions.
- We should test malformed-JSONPath behavior.
- One consistency note: end to end tests can be moved to
`VariantSuite.scala` to match similar variant expressions.
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -23416,6 +23416,45 @@ def is_variant_null(v: "ColumnOrName") -> Column:
return _invoke_function("is_variant_null", _to_java_column(v))
+@_try_remote_functions
+def variant_array_length(v: "ColumnOrName", path: Optional[Union[Column, str]]
= None) -> Column:
+ """
+ Returns the number of elements in the variant array at `path`. If `path`
is omitted, the root
+ array is inspected. Returns NULL if the input is SQL NULL, the path does
not exist, or the
+ target is a variant null or any non-array variant value.
+
+ .. versionadded:: 5.0.0
+
+ Parameters
+ ----------
+ v : :class:`~pyspark.sql.Column` or str
+ a variant column or column name
+ A column that evaluates to a variant.
+ path : :class:`~pyspark.sql.Column` or str, optional
+ the JSONPath identifying the array to inspect. A `str` is a literal
path; a
+ :class:`~pyspark.sql.Column` supplies the path at runtime. If omitted,
the root array is
+ inspected.
+
+ Returns
+ -------
+ :class:`~pyspark.sql.Column`
+ an integer column representing the array length, or NULL for non-array
variant values
+ Returns a column that evaluates to an integer.
+
+ Examples
+ --------
+ >>> df = spark.createDataFrame([('''{"a": [1, 2, 3]}''',), ('''{"a":
1}''',)], ['json'])
+ >>> df.select(variant_array_length(parse_json(df.json),
"$.a").alias("r")).collect()
+ [Row(r=3), Row(r=None)]
+ """
Review Comment:
We should add one example where we omit the path argument.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -123,6 +123,165 @@ case class IsVariantNull(child: Expression) extends
UnaryExpression
copy(child = newChild)
}
+case class VariantArrayLength(child: Expression)
+ extends UnaryExpression
+ with ExpectsInputTypes
+ with RuntimeReplaceable {
+
+ override lazy val replacement: Expression = StaticInvoke(
+ VariantArrayLength.getClass,
+ IntegerType,
+ "variantArrayLength",
+ Seq(child),
+ inputTypes,
+ returnNullable = true)
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(VariantType)
+
+ override def dataType: DataType = IntegerType
+
+ override def nullable: Boolean = true
+
+ override def prettyName: String = "variant_array_length"
+
+ override protected def withNewChildInternal(newChild: Expression):
VariantArrayLength =
+ copy(child = newChild)
+}
+
+case class VariantArrayLengthWithPath(child: Expression, path: Expression)
+ extends BinaryExpression
+ with ExpectsInputTypes {
+
+ @transient private lazy val parsedPath: Option[Array[VariantPathSegment]] = {
+ if (path.foldable) {
+ Option(path.eval()).map(p => VariantGet.getParsedPath(p.toString,
prettyName))
+ } else {
+ None
+ }
+ }
+
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(VariantType, StringTypeWithCollation(supportsTrimCollation = true))
+
+ override def dataType: DataType = IntegerType
+
+ override def nullable: Boolean = true
+
+ override def nullIntolerant: Boolean = true
+
+ override def prettyName: String = "variant_array_length"
+
+ override def eval(input: InternalRow): Any = {
+ val _ = parsedPath
+ super.eval(input)
+ }
+
+ override protected def nullSafeEval(input: Any, path: Any): Any = parsedPath
match {
+ case Some(pp) =>
+ VariantArrayLength.variantArrayLength(input.asInstanceOf[VariantVal], pp)
+ case _ =>
+ VariantArrayLength.variantArrayLength(
+ input.asInstanceOf[VariantVal], path.asInstanceOf[UTF8String],
prettyName)
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val childCode = child.genCode(ctx)
+ val (pathCode, pathArg, functionNameArg) = if (parsedPath.isEmpty) {
+ val pathCode = path.genCode(ctx)
+ (pathCode, pathCode.value, s""", "$prettyName"""")
+ } else {
+ (
+ new ExprCode(EmptyBlock, FalseLiteral, TrueLiteral),
+ ctx.addReferenceObj("parsedPath", parsedPath.get),
+ ""
+ )
+ }
+ val code = code"""
+ ${childCode.code}
+ ${pathCode.code}
+ boolean ${ev.isNull} = ${childCode.isNull} || ${pathCode.isNull};
+ int ${ev.value} = ${CodeGenerator.defaultValue(IntegerType)};
+ if (!${ev.isNull}) {
+ Integer length =
+
org.apache.spark.sql.catalyst.expressions.variant.VariantArrayLength.variantArrayLength(
+ ${childCode.value}, $pathArg$functionNameArg);
+ if (length == null) {
+ ${ev.isNull} = true;
+ } else {
+ ${ev.value} = length;
+ }
+ }
+ """
+ ev.copy(code = code)
+ }
+
+ override def left: Expression = child
+
+ override def right: Expression = path
+
+ override protected def withNewChildrenInternal(
+ newChild: Expression,
+ newPath: Expression): VariantArrayLengthWithPath = copy(child =
newChild, path = newPath)
+}
+
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(expr[, path]) - Returns the number of elements in the
variant array at `path`. " +
+ "If `path` is omitted, the root array is inspected. Returns NULL if the
input is SQL NULL, " +
+ "the path does not exist, or the target is a variant null or any non-array
variant value.",
+ arguments = """
+ Arguments:
+ * expr - A variant value to inspect.
+ * path - An optional string expression in JSONPath format that
identifies the array to
+ inspect. When omitted, it defaults to `$`.
+ """,
+ examples = """
+ Examples:
+ > SELECT _FUNC_(parse_json('[1, 2, 3]'));
+ 3
+ > SELECT _FUNC_(parse_json('{"a": [1, 2]}'), '$.a');
+ 2
+ > SELECT _FUNC_(parse_json('{"a": 1}'));
+ NULL
+ > SELECT _FUNC_(parse_json('null'));
Review Comment:
Let's add SQL NULL example too.
##########
sql/core/src/test/resources/sql-functions/sql-expression-schema.md:
##########
@@ -591,4 +592,4 @@
| org.apache.spark.sql.catalyst.expressions.xml.XPathList | xpath | SELECT
xpath('<a><b>b1</b><b>b2</b><b>b3</b><c>c1</c><c>c2</c></a>','a/b/text()') |
struct<xpath(<a><b>b1</b><b>b2</b><b>b3</b><c>c1</c><c>c2</c></a>,
a/b/text()):array<string>> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathLong | xpath_long |
SELECT xpath_long('<a><b>1</b><b>2</b></a>', 'sum(a/b)') |
struct<xpath_long(<a><b>1</b><b>2</b></a>, sum(a/b)):bigint> |
| org.apache.spark.sql.catalyst.expressions.xml.XPathShort | xpath_short |
SELECT xpath_short('<a><b>1</b><b>2</b></a>', 'sum(a/b)') |
struct<xpath_short(<a><b>1</b><b>2</b></a>, sum(a/b)):smallint> |
-| org.apache.spark.sql.catalyst.expressions.xml.XPathString | xpath_string |
SELECT xpath_string('<a><b>b</b><c>cc</c></a>','a/c') |
struct<xpath_string(<a><b>b</b><c>cc</c></a>, a/c):string> |
+| org.apache.spark.sql.catalyst.expressions.xml.XPathString | xpath_string |
SELECT xpath_string('<a><b>b</b><c>cc</c></a>','a/c') |
struct<xpath_string(<a><b>b</b><c>cc</c></a>, a/c):string> |
Review Comment:
Let's revert this as it is unrelated to this PR.
##########
python/pyspark/sql/functions/builtin.py:
##########
@@ -23416,6 +23416,45 @@ def is_variant_null(v: "ColumnOrName") -> Column:
return _invoke_function("is_variant_null", _to_java_column(v))
+@_try_remote_functions
+def variant_array_length(v: "ColumnOrName", path: Optional[Union[Column, str]]
= None) -> Column:
+ """
+ Returns the number of elements in the variant array at `path`. If `path`
is omitted, the root
+ array is inspected. Returns NULL if the input is SQL NULL, the path does
not exist, or the
+ target is a variant null or any non-array variant value.
Review Comment:
```suggestion
Returns the number of elements in the variant array at `path`. If `path`
is omitted, the root
array is inspected. Returns NULL if the input is SQL NULL, the path does
not exist, or the
target is not an array.
```
Saying "or the target is a variant null" is slightly redundant since it is
implied by the following statement.
Same for docs below.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -549,6 +708,18 @@ case object VariantGet {
parsedPath: Array[VariantPathSegment],
dataType: DataType,
castArgs: VariantCastArgs): Any = {
+ val v = getVariant(input, parsedPath)
+ if (v == null) {
+ null
+ } else {
+ VariantGet.cast(v, dataType, castArgs)
+ }
+ }
+
+ /**
+ * Returns the sub-variant at `parsedPath` without copying its value or
metadata.
+ */
+ def getVariant(input: VariantVal, parsedPath: Array[VariantPathSegment]):
Variant = {
Review Comment:
Having both `variantGet` and `getVariant` is easy to confuse, and the key
distinction (that this one returns an uncopied sub-variant sharing the parent's
metadata) is only conveyed by the comment. Could we rename it to something else?
--
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]
