chenhao-db commented on code in PR #45806:
URL: https://github.com/apache/spark/pull/45806#discussion_r1550297032
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/variant/variantExpressions.scala:
##########
@@ -403,3 +405,134 @@ object VariantGetExpressionBuilder extends
VariantGetExpressionBuilderBase(true)
)
// scalastyle:on line.size.limit
object TryVariantGetExpressionBuilder extends
VariantGetExpressionBuilderBase(false)
+
+@ExpressionDescription(
+ usage = "_FUNC_(v) - Returns schema in the SQL format of a variant.",
+ examples = """
+ Examples:
+ > SELECT _FUNC_(parse_json('null'));
+ VOID
+ > SELECT _FUNC_(parse_json('[{"b":true,"a":0}]'));
+ ARRAY<STRUCT<a: BIGINT, b: BOOLEAN>>
+ """,
+ since = "4.0.0",
+ group = "variant_funcs"
+)
+case class SchemaOfVariant(child: Expression)
+ extends UnaryExpression
+ with RuntimeReplaceable
+ with ExpectsInputTypes {
+ override lazy val replacement: Expression = StaticInvoke(
+ SchemaOfVariant.getClass,
+ StringType,
+ "schemaOfVariant",
+ Seq(child),
+ inputTypes,
+ returnNullable = false)
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(VariantType)
+
+ override def dataType: DataType = StringType
+
+ override def prettyName: String = "schema_of_variant"
+
+ override protected def withNewChildInternal(newChild: Expression):
SchemaOfVariant =
+ copy(child = newChild)
+}
+
+object SchemaOfVariant {
+ /** The actual implementation of the `SchemaOfVariant` expression. */
+ def schemaOfVariant(input: VariantVal): UTF8String = {
+ val v = new Variant(input.getValue, input.getMetadata)
+ UTF8String.fromString(schemaOf(v).sql)
+ }
+
+ /**
+ * Return the schema of a variant. Struct fields are guaranteed to be sorted
alphabetically.
+ */
+ def schemaOf(v: Variant): DataType = v.getType match {
+ case Type.OBJECT =>
+ val size = v.objectSize()
+ val fields = new Array[StructField](size)
+ for (i <- 0 until size) {
+ val field = v.getFieldAtIndex(i)
+ fields(i) = StructField(field.key, schemaOf(field.value))
+ }
+ // According to the variant spec, object fields must be sorted
alphabetically. So we don't
+ // have to sort, but just need to validate they are sorted.
+ for (i <- 1 until size) {
+ if (fields(i - 1).name >= fields(i).name) {
+ throw new SparkRuntimeException("MALFORMED_VARIANT", Map.empty)
+ }
+ }
+ StructType(fields)
+ case Type.ARRAY =>
+ var elementType: DataType = NullType
+ for (i <- 0 until v.arraySize()) {
+ elementType = mergeSchema(elementType,
schemaOf(v.getElementAtIndex(i)))
+ }
+ ArrayType(elementType)
+ case Type.NULL => NullType
+ case Type.BOOLEAN => BooleanType
+ case Type.LONG => LongType
+ case Type.STRING => StringType
+ case Type.DOUBLE => DoubleType
+ case Type.DECIMAL =>
+ val d = v.getDecimal
+ DecimalType(d.precision(), d.scale())
+ }
+
+ /**
+ * Returns the tightest common type for two given data types. Input struct
fields are assumed to
+ * be sorted alphabetically.
+ */
+ def mergeSchema(t1: DataType, t2: DataType): DataType = (t1, t2) match {
+ case (t1, t2) if t1 == t2 => t1
Review Comment:
I personally like the current code. I don't think `TypeCoercion` contains a
suitable rule that can be used here. If we use `findTightestCommonType`, we
still need most of the code that handles decimal/struct/array and can hardly
simplify the code. If we use any function that calls `findWiderTypeForDecimal`
(like `findWiderTypeForTwo`), its semantics will be undesired because If the
wider decimal type exceeds system limitation, this rule will truncate the
decimal type (and we still need custom code for struct/array). Using these
rules may fruitlessly visit the whole type object, and we need to do a second
pass of visit. Since this function is used in the expression evaluation, I
think we do care about its efficiency.
Essentially, `mergeSchema` only need to handle the result of `mergeSchema`
and `schemaOf`, and we can have a better control over them if we have all the
type resolution logic inside and avoid calling any libaray functions.
--
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]
