hudi-core 0.6.0

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use crate::error::Result;
use apache_avro::Schema as AvroSchema;
use apache_avro::schema::{Alias, DecimalSchema, EnumSchema, FixedSchema, Name, RecordSchema};
use apache_avro::types::Value;
use arrow::datatypes::{DataType, IntervalUnit, Schema, TimeUnit, UnionMode};
use arrow::datatypes::{Field, UnionFields};
use std::collections::HashMap;
use std::sync::Arc;

/// Timezone an Avro `timestamp-*` logical type denotes, per the Avro spec.
const UTC_TIMEZONE: &str = "UTC";

/// Converts an avro schema to an arrow schema
pub fn to_arrow_schema(avro_schema: &apache_avro::Schema) -> Result<Schema> {
    let mut schema_fields = vec![];
    match avro_schema {
        AvroSchema::Record(RecordSchema { fields, .. }) => {
            for field in fields {
                schema_fields.push(schema_to_field_with_props(
                    &field.schema,
                    Some(&field.name),
                    field.is_nullable(),
                    Some(external_props(&field.schema)),
                )?)
            }
        }
        schema => schema_fields.push(schema_to_field(schema, Some(""), false)?),
    }

    let schema = Schema::new(schema_fields);
    Ok(schema)
}

fn schema_to_field(
    schema: &apache_avro::Schema,
    name: Option<&str>,
    nullable: bool,
) -> Result<Field> {
    schema_to_field_with_props(schema, name, nullable, Default::default())
}

/// Arrow's conventional name for a map's entry struct, and what the parquet
/// reader produces — the two schemas have to agree by name to reconcile.
pub(crate) const MAP_ENTRIES_FIELD: &str = "key_value";

fn schema_to_field_with_props(
    schema: &AvroSchema,
    name: Option<&str>,
    nullable: bool,
    props: Option<HashMap<String, String>>,
) -> Result<Field> {
    let mut nullable = nullable;
    let field_type: DataType = match schema {
        AvroSchema::Ref { .. } => todo!("Add support for AvroSchema::Ref"),
        AvroSchema::Null => DataType::Null,
        AvroSchema::Boolean => DataType::Boolean,
        AvroSchema::Int => DataType::Int32,
        AvroSchema::Long => DataType::Int64,
        AvroSchema::Float => DataType::Float32,
        AvroSchema::Double => DataType::Float64,
        AvroSchema::Bytes => DataType::Binary,
        AvroSchema::String => DataType::Utf8,
        AvroSchema::Array(item_schema) => DataType::List(Arc::new(schema_to_field_with_props(
            &item_schema.items,
            Some("element"),
            false,
            None,
        )?)),
        AvroSchema::Map(value_schema) => {
            // An Arrow dictionary key must be an integer, so the `Dictionary(Utf8, V)`
            // this used to produce was not a valid type and could not reconcile
            // against the `Map` a parquet base file carries. Avro map keys are
            // always strings, and the entry struct is never null.
            let value_field =
                schema_to_field_with_props(&value_schema.types, Some("value"), true, None)?;
            let entries = Field::new(
                MAP_ENTRIES_FIELD,
                DataType::Struct(
                    vec![Field::new("key", DataType::Utf8, false), value_field].into(),
                ),
                false,
            );
            DataType::Map(Arc::new(entries), false)
        }
        AvroSchema::Union(us) => {
            // If there are only two variants and one of them is null, set the other type as the field data type
            let has_nullable = us
                .find_schema_with_known_schemata::<apache_avro::Schema>(&Value::Null, None, &None)
                .is_some();
            let sub_schemas = us.variants();
            if has_nullable && sub_schemas.len() == 2 {
                nullable = true;
                if let Some(schema) = sub_schemas
                    .iter()
                    .find(|&schema| !matches!(schema, AvroSchema::Null))
                {
                    schema_to_field_with_props(schema, None, has_nullable, None)?
                        .data_type()
                        .clone()
                } else {
                    return Err(apache_avro::Error::new(
                        apache_avro::error::Details::GetUnionDuplicate,
                    )
                    .into());
                }
            } else {
                let fields = sub_schemas
                    .iter()
                    .map(|s| schema_to_field_with_props(s, None, has_nullable, None))
                    .collect::<Result<Vec<Field>>>()?;
                let type_ids = 0_i8..fields.len() as i8;
                DataType::Union(UnionFields::try_new(type_ids, fields)?, UnionMode::Dense)
            }
        }
        AvroSchema::Record(RecordSchema { fields, .. }) => {
            let fields: Result<_> = fields
                .iter()
                .map(|field| {
                    let mut props = HashMap::new();
                    if let Some(doc) = &field.doc {
                        props.insert("avro::doc".to_string(), doc.clone());
                    }
                    /*if let Some(aliases) = fields.aliases {
                        props.insert("aliases", aliases);
                    }*/
                    schema_to_field_with_props(&field.schema, Some(&field.name), false, Some(props))
                })
                .collect();
            DataType::Struct(fields?)
        }
        AvroSchema::Enum(EnumSchema { .. }) => DataType::Utf8,
        AvroSchema::Fixed(FixedSchema { size, .. }) => DataType::FixedSizeBinary(*size as i32),
        AvroSchema::Decimal(DecimalSchema {
            precision, scale, ..
        }) => DataType::Decimal128(*precision as u8, *scale as i8),
        AvroSchema::BigDecimal => DataType::LargeBinary,
        AvroSchema::Uuid => DataType::FixedSizeBinary(16),
        AvroSchema::Date => DataType::Date32,
        AvroSchema::TimeMillis => DataType::Time32(TimeUnit::Millisecond),
        AvroSchema::TimeMicros => DataType::Time64(TimeUnit::Microsecond),
        // Avro's `timestamp-*` logical types are instants in UTC; the
        // timezone-naive variants are `local-timestamp-*` below. Carrying the
        // zone matters beyond correctness: a parquet base file reads back as
        // `Timestamp(_, "UTC")`, so dropping it here makes a log batch fail to
        // concatenate with the base batch it is merged against.
        AvroSchema::TimestampMillis => {
            DataType::Timestamp(TimeUnit::Millisecond, Some(UTC_TIMEZONE.into()))
        }
        AvroSchema::TimestampMicros => {
            DataType::Timestamp(TimeUnit::Microsecond, Some(UTC_TIMEZONE.into()))
        }
        AvroSchema::TimestampNanos => {
            DataType::Timestamp(TimeUnit::Nanosecond, Some(UTC_TIMEZONE.into()))
        }
        // Timezone-naive: no zone, matching arrow-avro's decode of these types and
        // a parquet base file written with isAdjustedToUTC=false.
        AvroSchema::LocalTimestampMillis => DataType::Timestamp(TimeUnit::Millisecond, None),
        AvroSchema::LocalTimestampMicros => DataType::Timestamp(TimeUnit::Microsecond, None),
        AvroSchema::LocalTimestampNanos => DataType::Timestamp(TimeUnit::Nanosecond, None),
        AvroSchema::Duration => DataType::Duration(TimeUnit::Millisecond),
    };

    let data_type = field_type.clone();
    let name = name.unwrap_or_else(|| default_field_name(&data_type));

    let mut field = Field::new(name, field_type, nullable);
    field.set_metadata(props.unwrap_or_default());
    Ok(field)
}

fn default_field_name(dt: &DataType) -> &str {
    match dt {
        DataType::Null => "null",
        DataType::Boolean => "bit",
        DataType::Int8 => "tinyint",
        DataType::Int16 => "smallint",
        DataType::Int32 => "int",
        DataType::Int64 => "bigint",
        DataType::UInt8 => "uint1",
        DataType::UInt16 => "uint2",
        DataType::UInt32 => "uint4",
        DataType::UInt64 => "uint8",
        DataType::Float16 => "float2",
        DataType::Float32 => "float4",
        DataType::Float64 => "float8",
        DataType::Date32 => "dateday",
        DataType::Date64 => "datemilli",
        DataType::Time32(tu) | DataType::Time64(tu) => match tu {
            TimeUnit::Second => "timesec",
            TimeUnit::Millisecond => "timemilli",
            TimeUnit::Microsecond => "timemicro",
            TimeUnit::Nanosecond => "timenano",
        },
        DataType::Timestamp(tu, tz) => {
            if tz.is_some() {
                match tu {
                    TimeUnit::Second => "timestampsectz",
                    TimeUnit::Millisecond => "timestampmillitz",
                    TimeUnit::Microsecond => "timestampmicrotz",
                    TimeUnit::Nanosecond => "timestampnanotz",
                }
            } else {
                match tu {
                    TimeUnit::Second => "timestampsec",
                    TimeUnit::Millisecond => "timestampmilli",
                    TimeUnit::Microsecond => "timestampmicro",
                    TimeUnit::Nanosecond => "timestampnano",
                }
            }
        }
        DataType::Duration(_) => "duration",
        DataType::Interval(unit) => match unit {
            IntervalUnit::YearMonth => "intervalyear",
            IntervalUnit::DayTime => "intervalmonth",
            IntervalUnit::MonthDayNano => "intervalmonthdaynano",
        },
        DataType::Binary => "varbinary",
        DataType::FixedSizeBinary(_) => "fixedsizebinary",
        DataType::LargeBinary => "largevarbinary",
        DataType::Utf8 => "varchar",
        DataType::LargeUtf8 => "largevarchar",
        DataType::List(_) => "list",
        DataType::FixedSizeList(_, _) => "fixed_size_list",
        DataType::LargeList(_) => "largelist",
        DataType::Struct(_) => "struct",
        DataType::Union(_, _) => "union",
        DataType::Dictionary(_, _) => "dictionary",
        DataType::Map(_, _) => "map",
        DataType::RunEndEncoded(_, _) => {
            unimplemented!("RunEndEncoded support not implemented")
        }
        DataType::Utf8View
        | DataType::BinaryView
        | DataType::ListView(_)
        | DataType::LargeListView(_) => {
            unimplemented!("View support not implemented")
        }
        DataType::Decimal32(_, _) => "decimal",
        DataType::Decimal64(_, _) => "decimal",
        DataType::Decimal128(_, _) => "decimal",
        DataType::Decimal256(_, _) => "decimal",
    }
}

fn external_props(schema: &AvroSchema) -> HashMap<String, String> {
    let mut props = HashMap::new();
    match &schema {
        AvroSchema::Record(RecordSchema { doc: Some(doc), .. })
        | AvroSchema::Enum(EnumSchema { doc: Some(doc), .. })
        | AvroSchema::Fixed(FixedSchema { doc: Some(doc), .. }) => {
            props.insert("avro::doc".to_string(), doc.clone());
        }
        _ => {}
    }
    match &schema {
        AvroSchema::Record(RecordSchema {
            name: Name { namespace, .. },
            aliases: Some(aliases),
            ..
        })
        | AvroSchema::Enum(EnumSchema {
            name: Name { namespace, .. },
            aliases: Some(aliases),
            ..
        })
        | AvroSchema::Fixed(FixedSchema {
            name: Name { namespace, .. },
            aliases: Some(aliases),
            ..
        }) => {
            let aliases: Vec<String> = aliases
                .iter()
                .map(|alias| aliased(alias, namespace.as_deref(), None))
                .collect();
            props.insert(
                "avro::aliases".to_string(),
                format!("[{}]", aliases.join(",")),
            );
        }
        _ => {}
    }
    props
}

/// Returns the fully qualified name for a field
pub fn aliased(alias: &Alias, namespace: Option<&str>, default_namespace: Option<&str>) -> String {
    if alias.namespace().is_some() {
        alias.fullname(None)
    } else {
        let namespace = namespace.as_ref().copied().or(default_namespace);

        match namespace {
            Some(ref namespace) => format!("{}.{}", namespace, alias.name()),
            None => alias.fullname(None),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use apache_avro::Schema as AvroSchema;

    /// The timezone-naive `local-timestamp-*` types map to a tz-less Arrow timestamp,
    /// while the UTC `timestamp-*` types keep their zone.
    #[test]
    fn test_local_timestamps_convert_to_tz_less_timestamps() {
        let avro = AvroSchema::parse_str(
            r#"{"type":"record","name":"r","fields":[
                 {"name":"l_ms","type":{"type":"long","logicalType":"local-timestamp-millis"}},
                 {"name":"l_us","type":{"type":"long","logicalType":"local-timestamp-micros"}},
                 {"name":"l_ns","type":{"type":"long","logicalType":"local-timestamp-nanos"}},
                 {"name":"u_us","type":{"type":"long","logicalType":"timestamp-micros"}},
                 {"name":"opt","type":["null",{"type":"long","logicalType":"local-timestamp-micros"}]}]}"#,
        )
        .unwrap();

        let schema = to_arrow_schema(&avro).unwrap();
        let types: Vec<&DataType> = schema.fields().iter().map(|f| f.data_type()).collect();
        assert_eq!(
            types,
            vec![
                &DataType::Timestamp(TimeUnit::Millisecond, None),
                &DataType::Timestamp(TimeUnit::Microsecond, None),
                &DataType::Timestamp(TimeUnit::Nanosecond, None),
                &DataType::Timestamp(TimeUnit::Microsecond, Some(UTC_TIMEZONE.into())),
                &DataType::Timestamp(TimeUnit::Microsecond, None),
            ]
        );
        assert!(schema.field_with_name("opt").unwrap().is_nullable());
    }

    /// An Avro map becomes an Arrow `Map`, not a dictionary. A dictionary key
    /// must be an integer, so the `Dictionary(Utf8, V)` this used to produce was
    /// not a valid Arrow type and could not reconcile against the `Map` a
    /// parquet base file carries.
    #[test]
    fn test_avro_map_converts_to_arrow_map() {
        let avro = AvroSchema::parse_str(
            r#"{"type":"record","name":"r","fields":[
                 {"name":"m","type":{"type":"map","values":"int"}}]}"#,
        )
        .unwrap();

        let schema = to_arrow_schema(&avro).unwrap();
        let DataType::Map(entries, sorted) = schema.field(0).data_type() else {
            panic!("expected a Map, got {}", schema.field(0).data_type());
        };
        assert!(!sorted);
        assert_eq!(entries.name(), MAP_ENTRIES_FIELD);
        assert!(!entries.is_nullable(), "map entries are never null");

        let DataType::Struct(fields) = entries.data_type() else {
            panic!("entries must be a struct");
        };
        assert_eq!(fields.len(), 2);
        assert_eq!(fields[0].name(), "key");
        assert_eq!(fields[0].data_type(), &DataType::Utf8);
        assert!(!fields[0].is_nullable(), "an Avro map key is never null");
        assert_eq!(fields[1].name(), "value");
        assert_eq!(fields[1].data_type(), &DataType::Int32);
    }

    /// A map whose values are records nests as a struct inside the entry, rather
    /// than collapsing to the value type alone.
    #[test]
    fn test_avro_map_of_records_nests_the_value_struct() {
        let avro = AvroSchema::parse_str(
            r#"{"type":"record","name":"r","fields":[
                 {"name":"m","type":{"type":"map","values":
                   {"type":"record","name":"v","fields":[{"name":"a","type":"double"}]}}}]}"#,
        )
        .unwrap();

        let schema = to_arrow_schema(&avro).unwrap();
        let DataType::Map(entries, _) = schema.field(0).data_type() else {
            panic!("expected a Map");
        };
        let DataType::Struct(fields) = entries.data_type() else {
            panic!("entries must be a struct");
        };
        let DataType::Struct(value_fields) = fields[1].data_type() else {
            panic!(
                "map value must stay a struct, got {}",
                fields[1].data_type()
            );
        };
        assert_eq!(value_fields[0].name(), "a");
        assert_eq!(value_fields[0].data_type(), &DataType::Float64);
    }
}