uros-b commented on code in PR #58301:
URL: https://github.com/apache/spark/pull/58301#discussion_r4070955358
##########
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'));
+ NULL
+ """,
+ since = "5.0.0",
Review Comment:
```suggestion
since = "4.4.0",
```
--
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]
