bojana-db commented on code in PR #57125:
URL: https://github.com/apache/spark/pull/57125#discussion_r3569420380
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -942,6 +942,177 @@ object VariantInsert {
}
}
+case class VariantSet(
+ input: Expression,
+ path: Expression,
+ value: Expression,
+ createIfMissing: Expression)
+ extends QuaternaryExpression
+ with ExpectsInputTypes
+ with QueryErrorsBase {
+
+ override def first: Expression = input
+ override def second: Expression = path
+ override def third: Expression = value
+ override def fourth: Expression = createIfMissing
+
+ override def nullIntolerant: Boolean = true
+
+ override def dataType: DataType = VariantType
+ override def inputTypes: Seq[AbstractDataType] =
+ Seq(
+ VariantType,
+ StringTypeWithCollation(supportsTrimCollation = true),
+ AnyDataType,
+ BooleanType)
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ val result = super.checkInputDataTypes()
+ if (result.isFailure) {
+ result
+ } else if (value.dataType == NullType) {
+ TypeCheckResult.TypeCheckSuccess
+ } else if (!VariantGet.checkDataType(value.dataType, allowStructsAndMaps =
false)) {
+ DataTypeMismatch(
+ errorSubClass = "CAST_WITHOUT_SUGGESTION",
+ messageParameters =
+ Map("srcType" -> toSQLType(value.dataType), "targetType" ->
toSQLType(VariantType)))
+ } else {
+ TypeCheckResult.TypeCheckSuccess
+ }
+ }
+
+ // When the path is a foldable expression, parse it once at planning time
and cache it. The Java
+ // segments are derived once per task (see `ParsedSetPath`), avoiding a
per-row conversion. `None`
+ // means the path is dynamic (or a foldable NULL, which makes the whole
expression NULL and is
+ // never evaluated).
+ @transient private lazy val foldablePath: Option[VariantSet.ParsedSetPath] =
{
+ if (path.foldable) {
+ val p = path.eval()
+ if (p == null) {
+ None
+ } else {
+ val s = p.asInstanceOf[UTF8String].toString
+ Some(VariantSet.ParsedSetPath(
+ VariantExpressionEvalUtils.parseVariantPath(s, prettyName), s))
+ }
+ } else {
+ None
+ }
+ }
+
+ override protected def nullSafeEval(v: Any, p: Any, valValue: Any, create:
Any): Any = {
+ val inputVariant = v.asInstanceOf[VariantVal]
+ val createIfMissingValue = create.asInstanceOf[Boolean]
+ foldablePath match {
+ case Some(parsed) =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, parsed.javaSegments, parsed.pathStr, valValue,
value.dataType,
+ createIfMissingValue, prettyName)
+ case None =>
+ VariantExpressionEvalUtils.setAtPath(
+ inputVariant, p.asInstanceOf[UTF8String], valValue, value.dataType,
+ createIfMissingValue, prettyName)
+ }
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ val cls = VariantExpressionEvalUtils.getClass.getName.stripSuffix("$")
+ nullSafeCodeGen(ctx, ev, (vVal, pVal, valVal, createVal) => {
+ val fromArg = ctx.addReferenceObj("from", value.dataType)
+ foldablePath match {
+ case Some(parsed) =>
+ val parsedArg = ctx.addReferenceObj("setPath", parsed)
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $parsedArg.javaSegments(), $parsedArg.pathStr(),
$valVal, $fromArg,
+ | $createVal, "$prettyName");
+ """.stripMargin
+ case None =>
+ s"""
+ |${ev.value} = $cls.setAtPath(
+ | $vVal, $pVal, $valVal, $fromArg, $createVal, "$prettyName");
+ """.stripMargin
+ }
+ })
+ }
+
+ override def prettyName: String = "variant_set"
+
+ override protected def withNewChildrenInternal(
+ newFirst: Expression,
+ newSecond: Expression,
+ newThird: Expression,
+ newFourth: Expression): VariantSet =
+ copy(input = newFirst, path = newSecond, value = newThird, createIfMissing
= newFourth)
+}
+
+object VariantSet {
+ // Caches a foldable path. `VariantBuilder.PathSegment` is not
`Serializable`, so the Java form is
+ // `@transient` and re-derived once per executor task after deserialization.
`pathStr` is the
+ // source string, retained for error messages.
+ case class ParsedSetPath(segments: Array[VariantPathSegment], pathStr:
String) {
+ @transient lazy val javaSegments: Array[VariantBuilder.PathSegment] =
+ VariantExpressionEvalUtils.toJavaSegments(segments)
+ }
Review Comment:
Yes, once the whole family is merged I'll do a dedicated refactor pass over
all of 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]
