JingsongLi commented on code in PR #459:
URL: https://github.com/apache/paimon-rust/pull/459#discussion_r3526244210


##########
crates/paimon/src/variant.rs:
##########
@@ -1694,14 +1913,1179 @@ fn java_string_cmp(left: &str, right: &str) -> 
std::cmp::Ordering {
     left.encode_utf16().cmp(right.encode_utf16())
 }
 
+#[derive(Clone, Copy, Debug)]
+pub(crate) struct VariantShreddingInferConfig {
+    pub(crate) max_schema_depth: usize,
+    pub(crate) min_field_cardinality_ratio: f64,
+}
+
+pub(crate) fn infer_variant_shredding_schema(
+    variants: &[GenericVariant],
+    config: &VariantShreddingInferConfig,
+    max_fields_remaining: &mut usize,
+) -> Result<DataType> {
+    let mut simple_schema = None;
+    for variant in variants {
+        let schema_of_row = schema_of_variant(variant.as_ref()?, 
config.max_schema_depth)?;
+        simple_schema = merge_inferred_schema(simple_schema, schema_of_row)?;
+    }
+
+    let min_cardinality =
+        ((variants.len() as f64) * config.min_field_cardinality_ratio).ceil() 
as u64;
+    finalize_inferred_schema(simple_schema, min_cardinality, 
max_fields_remaining)
+}
+
+fn schema_of_variant(variant: VariantRef<'_>, max_depth: usize) -> 
Result<Option<DataType>> {
+    Ok(match variant.kind()? {
+        VariantKind::Object => {
+            if max_depth == 0 {
+                return Ok(Some(inferred_variant_type()));
+            }
+
+            let object_fields = variant.object_fields()?;
+            let mut fields = Vec::with_capacity(object_fields.len());
+            for (idx, field) in object_fields.iter().enumerate() {
+                let field_type = schema_of_variant(field.value, max_depth - 1)?
+                    .unwrap_or_else(inferred_variant_type);
+                fields.push(
+                    DataField::new(idx as i32, field.key.clone(), field_type)
+                        .with_description(Some("1".to_string())),
+                );
+            }
+
+            for idx in 1..fields.len() {
+                if java_string_cmp(fields[idx - 1].name(), fields[idx].name())
+                    != std::cmp::Ordering::Less
+                {
+                    return data_invalid("Variant object fields must be sorted 
alphabetically");
+                }
+            }
+
+            Some(DataType::Row(RowType::new(fields)))
+        }
+        VariantKind::Array => {
+            if max_depth == 0 {
+                return Ok(Some(inferred_variant_type()));
+            }
+
+            let mut element_type = None;
+            for element in variant.array_elements()? {
+                element_type = merge_inferred_schema(
+                    element_type,
+                    schema_of_variant(element, max_depth - 1)?,
+                )?;
+            }
+            Some(DataType::Array(ArrayType::new(
+                element_type.unwrap_or_else(inferred_variant_type),
+            )))
+        }
+        VariantKind::Null => None,
+        VariantKind::Boolean => Some(DataType::Boolean(BooleanType::new())),
+        VariantKind::Long => {
+            let value = variant.get_long()?;
+            let precision = decimal_precision(value as i128) as u32;
+            if precision <= 18 {
+                Some(DataType::Decimal(DecimalType::new(precision, 0)?))
+            } else {
+                Some(DataType::BigInt(BigIntType::new()))
+            }
+        }
+        VariantKind::String => 
Some(DataType::VarChar(VarCharType::string_type())),
+        VariantKind::Double => Some(DataType::Double(DoubleType::new())),
+        VariantKind::Decimal => {
+            let decimal = variant.get_decimal()?;
+            if decimal.scale < 0 {
+                Some(inferred_variant_type())
+            } else {
+                let scale = decimal.scale as u32;
+                let mut precision = decimal.precision as u32;
+                if precision < scale {
+                    precision = scale;
+                }
+                if precision == 0 {
+                    precision = 1;
+                }
+                Some(DataType::Decimal(DecimalType::new(precision, scale)?))
+            }
+        }
+        VariantKind::Date => Some(DataType::Date(DateType::new())),
+        VariantKind::Timestamp => Some(DataType::LocalZonedTimestamp(
+            
LocalZonedTimestampType::new(LocalZonedTimestampType::DEFAULT_PRECISION)?,
+        )),
+        VariantKind::TimestampNtz => 
Some(DataType::Timestamp(TimestampType::new(
+            TimestampType::DEFAULT_PRECISION,
+        )?)),
+        VariantKind::Float => Some(DataType::Float(FloatType::new())),
+        VariantKind::Binary => Some(DataType::VarBinary(VarBinaryType::try_new(
+            true,
+            VarBinaryType::MAX_LENGTH,
+        )?)),
+        VariantKind::Uuid => Some(inferred_variant_type()),
+    })
+}
+
+fn merge_inferred_schema(
+    left: Option<DataType>,
+    right: Option<DataType>,
+) -> Result<Option<DataType>> {
+    match (left, right) {
+        (None, right) => Ok(right),
+        (left, None) => Ok(left),
+        (Some(left), Some(right)) => merge_inferred_types(left, 
right).map(Some),
+    }
+}
+
+fn merge_inferred_types(left: DataType, right: DataType) -> Result<DataType> {
+    Ok(match (&left, &right) {
+        (DataType::Decimal(left), DataType::Decimal(right)) => {
+            merge_inferred_decimals(left, right)?
+        }
+        (DataType::Decimal(decimal), DataType::BigInt(_))
+        | (DataType::BigInt(_), DataType::Decimal(decimal)) => {
+            merge_inferred_decimal_with_long(decimal)?
+        }
+        (DataType::Row(left), DataType::Row(right)) => {
+            DataType::Row(merge_inferred_row_types(left, right)?)
+        }
+        (DataType::Array(left), DataType::Array(right)) => {
+            let element_type = merge_inferred_schema(
+                Some(left.element_type().clone()),
+                Some(right.element_type().clone()),
+            )?
+            .unwrap_or_else(inferred_variant_type);
+            DataType::Array(ArrayType::new(element_type))
+        }
+        _ if left == right => left,
+        _ => inferred_variant_type(),
+    })
+}
+
+fn merge_inferred_decimals(left: &DecimalType, right: &DecimalType) -> 
Result<DataType> {
+    let scale = left.scale().max(right.scale());
+    let range = (left.precision() - left.scale()).max(right.precision() - 
right.scale());
+    if range + scale > DecimalType::MAX_PRECISION {
+        Ok(inferred_variant_type())
+    } else {
+        Ok(DataType::Decimal(DecimalType::new(range + scale, scale)?))
+    }
+}
+
+fn merge_inferred_decimal_with_long(decimal: &DecimalType) -> Result<DataType> 
{
+    if decimal.scale() == 0 && decimal.precision() <= 18 {
+        Ok(DataType::BigInt(BigIntType::new()))
+    } else {
+        merge_inferred_decimals(decimal, &DecimalType::new(19, 0)?)
+    }
+}
+
+fn merge_inferred_row_types(left: &RowType, right: &RowType) -> 
Result<RowType> {
+    let fields1 = left.fields();
+    let fields2 = right.fields();
+    let mut new_fields = Vec::new();
+    let mut f1_idx = 0;
+    let mut f2_idx = 0;
+    let mut next_field_id = 0;
+    const MAX_ROW_FIELD_SIZE: usize = 1000;
+
+    while f1_idx < fields1.len() && f2_idx < fields2.len() && new_fields.len() 
< MAX_ROW_FIELD_SIZE
+    {
+        let field1 = &fields1[f1_idx];
+        let field2 = &fields2[f2_idx];
+        match java_string_cmp(field1.name(), field2.name()) {
+            std::cmp::Ordering::Equal => {
+                let data_type =
+                    merge_inferred_types(field1.data_type().clone(), 
field2.data_type().clone())?;
+                let count = inferred_field_count(field1)? + 
inferred_field_count(field2)?;
+                new_fields.push(inferred_count_field(
+                    next_field_id,
+                    field1.name(),
+                    data_type,
+                    count,
+                ));
+                next_field_id += 1;
+                f1_idx += 1;
+                f2_idx += 1;
+            }
+            std::cmp::Ordering::Less => {
+                let count = inferred_field_count(field1)?;
+                new_fields.push(inferred_count_field(
+                    next_field_id,
+                    field1.name(),
+                    field1.data_type().clone(),
+                    count,
+                ));
+                next_field_id += 1;
+                f1_idx += 1;
+            }
+            std::cmp::Ordering::Greater => {
+                let count = inferred_field_count(field2)?;
+                new_fields.push(inferred_count_field(
+                    next_field_id,
+                    field2.name(),
+                    field2.data_type().clone(),
+                    count,
+                ));
+                next_field_id += 1;
+                f2_idx += 1;
+            }
+        }
+    }
+
+    while f1_idx < fields1.len() && new_fields.len() < MAX_ROW_FIELD_SIZE {
+        let field = &fields1[f1_idx];
+        let count = inferred_field_count(field)?;
+        new_fields.push(inferred_count_field(
+            next_field_id,
+            field.name(),
+            field.data_type().clone(),
+            count,
+        ));
+        next_field_id += 1;
+        f1_idx += 1;
+    }
+
+    while f2_idx < fields2.len() && new_fields.len() < MAX_ROW_FIELD_SIZE {
+        let field = &fields2[f2_idx];
+        let count = inferred_field_count(field)?;
+        new_fields.push(inferred_count_field(
+            next_field_id,
+            field.name(),
+            field.data_type().clone(),
+            count,
+        ));
+        next_field_id += 1;
+        f2_idx += 1;
+    }
+
+    Ok(RowType::new(new_fields))
+}
+
+fn inferred_field_count(field: &DataField) -> Result<u64> {
+    let Some(description) = field.description() else {
+        return data_invalid("Variant inferred field is missing count");
+    };
+    description.parse::<u64>().map_err(|e| Error::DataInvalid {
+        message: format!("Invalid Variant inferred field count: 
{description}"),
+        source: Some(Box::new(e)),
+    })
+}
+
+fn inferred_count_field(id: i32, name: &str, data_type: DataType, count: u64) 
-> DataField {
+    DataField::new(id, name.to_string(), 
data_type).with_description(Some(count.to_string()))
+}
+
+fn finalize_inferred_schema(
+    data_type: Option<DataType>,
+    min_cardinality: u64,
+    max_fields_remaining: &mut usize,
+) -> Result<DataType> {
+    if *max_fields_remaining == 0 {
+        return Ok(inferred_variant_type());
+    }
+    *max_fields_remaining -= 1;
+    if *max_fields_remaining == 0 {
+        return Ok(inferred_variant_type());
+    }
+
+    let Some(data_type) = data_type else {
+        return Ok(inferred_variant_type());
+    };
+    if matches!(data_type, DataType::Variant(_)) {
+        return Ok(inferred_variant_type());
+    }
+
+    Ok(match data_type {
+        DataType::Row(row_type) => {
+            let mut fields = Vec::new();
+            for field in row_type.fields() {
+                if inferred_field_count(field)? >= min_cardinality && 
*max_fields_remaining > 0 {
+                    let field_type = finalize_inferred_schema(
+                        Some(field.data_type().clone()),
+                        min_cardinality,
+                        max_fields_remaining,
+                    )?;
+                    fields.push(DataField::new(
+                        fields.len() as i32,
+                        field.name().to_string(),
+                        field_type,
+                    ));
+                }
+            }
+
+            if fields.is_empty() {
+                inferred_variant_type()
+            } else {
+                DataType::Row(RowType::new(fields))
+            }
+        }
+        DataType::Array(array_type) => {
+            let element_type = finalize_inferred_schema(
+                Some(array_type.element_type().clone()),
+                min_cardinality,
+                max_fields_remaining,
+            )?;
+            DataType::Array(ArrayType::new(element_type))
+        }
+        DataType::TinyInt(_) | DataType::SmallInt(_) | DataType::Int(_) | 
DataType::BigInt(_) => {
+            *max_fields_remaining = (*max_fields_remaining).saturating_sub(1);
+            DataType::BigInt(BigIntType::new())
+        }
+        DataType::Decimal(decimal) => {
+            *max_fields_remaining = (*max_fields_remaining).saturating_sub(1);
+            if decimal.precision() <= 18 && decimal.scale() == 0 {
+                DataType::BigInt(BigIntType::new())
+            } else if decimal.precision() <= 18 {
+                DataType::Decimal(DecimalType::new(18, decimal.scale())?)
+            } else {
+                DataType::Decimal(DecimalType::new(
+                    DecimalType::MAX_PRECISION,
+                    decimal.scale(),
+                )?)
+            }
+        }
+        other => {
+            *max_fields_remaining = (*max_fields_remaining).saturating_sub(1);
+            other
+        }
+    })
+}
+
+fn inferred_variant_type() -> DataType {
+    DataType::Variant(VariantType::new())
+}
+
+pub(crate) fn variant_shredding_type(data_type: &DataType) -> Result<DataType> 
{
+    Ok(DataType::Row(variant_shredding_row_type(
+        data_type, true, false,
+    )?))
+}
+
+fn variant_shredding_row_type(
+    data_type: &DataType,
+    is_top_level: bool,
+    is_object_field: bool,
+) -> Result<RowType> {
+    let mut fields = Vec::new();
+    if is_top_level {
+        fields.push(DataField::new(
+            0,
+            VARIANT_METADATA_FIELD_NAME.to_string(),
+            variant_binary_type(false)?,
+        ));
+    }
+
+    match data_type {
+        DataType::Array(array_type) => {
+            let element = DataType::Row(variant_shredding_row_type(
+                array_type.element_type(),
+                false,
+                false,
+            )?);
+            fields.push(DataField::new(
+                1,
+                VARIANT_VALUE_FIELD_NAME.to_string(),
+                variant_binary_type(true)?,
+            ));
+            fields.push(DataField::new(
+                2,
+                VARIANT_TYPED_VALUE_FIELD_NAME.to_string(),
+                
DataType::Array(ArrayType::with_nullable(data_type.is_nullable(), element)),
+            ));
+        }
+        DataType::Row(row_type) => {
+            let object_fields = row_type
+                .fields()
+                .iter()
+                .map(|field| {
+                    let child =
+                        
DataType::Row(variant_shredding_row_type(field.data_type(), false, true)?)
+                            .copy_with_nullable(false)?;
+                    Ok(DataField::new(field.id(), field.name().to_string(), 
child)
+                        
.with_description(field.description().map(ToString::to_string)))
+                })
+                .collect::<Result<Vec<_>>>()?;
+            fields.push(DataField::new(
+                1,
+                VARIANT_VALUE_FIELD_NAME.to_string(),
+                variant_binary_type(true)?,
+            ));
+            fields.push(DataField::new(
+                2,
+                VARIANT_TYPED_VALUE_FIELD_NAME.to_string(),
+                DataType::Row(RowType::new(object_fields)),
+            ));
+        }
+        DataType::Variant(_) => {
+            fields.push(DataField::new(
+                1,
+                VARIANT_VALUE_FIELD_NAME.to_string(),
+                variant_binary_type(is_object_field)?,
+            ));
+        }
+        dt if scalar_schema_for_type(dt)?.is_some() => {
+            fields.push(DataField::new(
+                1,
+                VARIANT_VALUE_FIELD_NAME.to_string(),
+                variant_binary_type(true)?,
+            ));
+            fields.push(DataField::new(
+                2,
+                VARIANT_TYPED_VALUE_FIELD_NAME.to_string(),
+                data_type.clone(),
+            ));
+        }
+        other => return invalid_variant_shredding_schema(format!("{other:?}")),
+    }
+
+    Ok(RowType::new(fields))
+}
+
+fn variant_binary_type(nullable: bool) -> Result<DataType> {
+    Ok(DataType::VarBinary(VarBinaryType::try_new(
+        nullable,
+        VarBinaryType::MAX_LENGTH,
+    )?))
+}
+
+pub(crate) fn build_variant_schema(row_type: &RowType) -> 
Result<VariantSchema> {
+    build_variant_schema_with_level(row_type, true)
+}
+
+fn build_variant_schema_with_level(row_type: &RowType, top_level: bool) -> 
Result<VariantSchema> {
+    let fields = row_type.fields();
+    if fields.is_empty() {
+        return invalid_variant_shredding_schema(format!("{row_type:?}"));
+    }
+
+    let mut seen = HashSet::new();
+    let mut typed_idx = None;
+    let mut variant_idx = None;
+    let mut top_level_metadata_idx = None;
+    let mut scalar_schema = None;
+    let mut object_schema = None;
+    let mut array_schema = None;
+
+    for (idx, field) in fields.iter().enumerate() {
+        if !seen.insert(field.name().to_string()) {
+            return invalid_variant_shredding_schema(format!("{row_type:?}"));
+        }
+
+        match field.name() {
+            VARIANT_TYPED_VALUE_FIELD_NAME => {
+                if typed_idx.is_some() {
+                    return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                }
+                typed_idx = Some(idx);
+                match field.data_type() {
+                    DataType::Row(inner) => {
+                        if inner.fields().is_empty() {
+                            return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                        }
+                        let mut object_names = HashSet::new();
+                        let mut object_fields = 
Vec::with_capacity(inner.fields().len());
+                        for child in inner.fields() {
+                            if !object_names.insert(child.name().to_string()) {
+                                return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                            }
+                            let DataType::Row(child_row) = child.data_type() 
else {
+                                return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                            };
+                            object_fields.push(VariantObjectField {
+                                field_name: child.name().to_string(),
+                                schema: 
build_variant_schema_with_level(child_row, false)?,
+                            });
+                        }
+                        object_schema = Some(object_fields);
+                    }
+                    DataType::Array(array_type) => {
+                        let DataType::Row(element_row) = 
array_type.element_type() else {
+                            return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                        };
+                        array_schema = 
Some(Box::new(build_variant_schema_with_level(
+                            element_row,
+                            false,
+                        )?));
+                    }
+                    dt => {
+                        scalar_schema = scalar_schema_for_type(dt)?;
+                        if scalar_schema.is_none() {
+                            return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                        }
+                    }
+                }
+            }
+            VARIANT_VALUE_FIELD_NAME => {
+                if variant_idx.is_some() || 
!is_varbinary_type(field.data_type()) {
+                    return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                }
+                variant_idx = Some(idx);
+            }
+            VARIANT_METADATA_FIELD_NAME => {
+                if top_level_metadata_idx.is_some() || 
!is_varbinary_type(field.data_type()) {
+                    return 
invalid_variant_shredding_schema(format!("{row_type:?}"));
+                }
+                top_level_metadata_idx = Some(idx);
+            }
+            _ => return 
invalid_variant_shredding_schema(format!("{row_type:?}")),
+        }
+    }
+
+    if top_level != top_level_metadata_idx.is_some() {
+        return invalid_variant_shredding_schema(format!("{row_type:?}"));
+    }
+
+    let object_schema_map = object_schema
+        .as_ref()
+        .map(|fields| {
+            fields
+                .iter()
+                .enumerate()
+                .map(|(idx, field)| (field.field_name.clone(), idx))
+                .collect()
+        })
+        .unwrap_or_default();
+
+    Ok(VariantSchema {
+        typed_idx,
+        variant_idx,
+        top_level_metadata_idx,
+        num_fields: fields.len(),
+        scalar_schema,
+        object_schema,
+        object_schema_map,
+        array_schema,
+    })
+}
+
+fn is_varbinary_type(data_type: &DataType) -> bool {
+    matches!(data_type, DataType::VarBinary(_) | DataType::Binary(_))
+}
+
+fn scalar_schema_for_type(data_type: &DataType) -> 
Result<Option<VariantScalarSchema>> {
+    Ok(match data_type {
+        DataType::Boolean(_) => Some(VariantScalarSchema::Boolean),
+        DataType::TinyInt(_) => Some(VariantScalarSchema::Int8),
+        DataType::SmallInt(_) => Some(VariantScalarSchema::Int16),
+        DataType::Int(_) => Some(VariantScalarSchema::Int32),
+        DataType::BigInt(_) => Some(VariantScalarSchema::Int64),
+        DataType::Float(_) => Some(VariantScalarSchema::Float32),
+        DataType::Double(_) => Some(VariantScalarSchema::Float64),
+        DataType::Decimal(decimal) => Some(VariantScalarSchema::Decimal {
+            precision: u8::try_from(decimal.precision()).map_err(|_| 
Error::DataTypeInvalid {
+                message: "Variant shredding decimal precision exceeds 
u8".to_string(),
+            })?,
+            scale: i8::try_from(decimal.scale() as i32).map_err(|_| 
Error::DataTypeInvalid {
+                message: "Variant shredding decimal scale exceeds 
i8".to_string(),
+            })?,
+        }),
+        DataType::VarChar(_) | DataType::Char(_) => 
Some(VariantScalarSchema::String),
+        DataType::VarBinary(_) | DataType::Binary(_) | DataType::Blob(_) => {
+            Some(VariantScalarSchema::Binary)
+        }
+        DataType::Date(_) => Some(VariantScalarSchema::Date32),
+        DataType::Timestamp(_) | DataType::LocalZonedTimestamp(_) => {

Review Comment:
   Fixed. I split the scalar schema into `Timestamp` and `TimestampNtz`: 
`DataType::LocalZonedTimestamp` now maps to the Variant TIMESTAMP kind, while 
`DataType::Timestamp` maps to TIMESTAMP_NTZ. Typed shredding now only accepts 
the matching Variant kind, and rebuild calls `append_timestamp_ntz` for NTZ 
values. Added `cast_and_rebuild_shredded_timestamp_ntz_preserves_kind` to cover 
the reported case.



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

Reply via email to