diff --git a/Cargo.toml b/Cargo.toml index 23d2a5785..63e00c886 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,7 +49,7 @@ arrow-row = { version = "57" } arrow-schema = { version = "57", features = ["serde"] } arrow-select = { version = "57" } object_store = { version = "0.12", features = ["aws", "azure", "gcp"] } -parquet = { version = "57", features = ["async", "object_store"] } +parquet = { version = "57", features = ["async", "object_store", "variant_experimental"] } # avro apache-avro = { version = "0.21", features = ["derive"] } diff --git a/crates/core/src/avro_to_arrow/arrow_array_reader.rs b/crates/core/src/avro_to_arrow/arrow_array_reader.rs index eed1edd73..41c677336 100644 --- a/crates/core/src/avro_to_arrow/arrow_array_reader.rs +++ b/crates/core/src/avro_to_arrow/arrow_array_reader.rs @@ -930,3 +930,102 @@ where } } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::StructArray; + use parquet::variant::{Variant, VariantArray, VariantArrayBuilder, VariantBuilderExt}; + + fn variant_bytes() -> (Vec, Vec) { + let mut builder = VariantArrayBuilder::new(1); + builder.append_value("iceberg"); + let array = builder.build(); + let metadata = array.metadata_field().value(0).to_vec(); + let value = array.value_field().unwrap().value(0).to_vec(); + (metadata, value) + } + + fn variant_schema(child_fields: &str) -> AvroSchema { + AvroSchema::parse_str(&format!( + r#"{{ + "type": "record", + "name": "root", + "fields": [ + {{ + "name": "var", + "type": {{ + "type": "record", + "name": "variant_record", + "logicalType": "variant", + "fields": [{child_fields}] + }} + }} + ] + }}"# + )) + .unwrap() + } + + fn variant_value(fields: Vec<(String, Value)>) -> Value { + Value::Record(vec![("var".to_string(), Value::Record(fields))]) + } + + #[test] + fn reads_hudi_variant_avro_record_as_variant_array() { + let schema = variant_schema( + r#" + {"name": "metadata", "type": "bytes"}, + {"name": "value", "type": "bytes"} + "#, + ); + let (metadata, value) = variant_bytes(); + let value = variant_value(vec![ + ("metadata".to_string(), Value::Bytes(metadata)), + ("value".to_string(), Value::Bytes(value)), + ]); + + let mut reader = AvroArrowArrayReader::try_new([Ok(value)].into_iter(), &schema).unwrap(); + let batch = reader.next_batch(10).unwrap().unwrap(); + let schema = batch.schema(); + let field = schema.field_with_name("var").unwrap(); + + field + .try_extension_type::() + .unwrap(); + let variant_array = VariantArray::try_new(batch.column_by_name("var").unwrap()).unwrap(); + assert_eq!(variant_array.value(0), Variant::from("iceberg")); + } + + #[test] + fn reads_hudi_variant_avro_record_by_child_name_not_position() { + let schema = variant_schema( + r#" + {"name": "value", "type": "bytes"}, + {"name": "metadata", "type": "bytes"} + "#, + ); + let (metadata, value) = variant_bytes(); + let value = variant_value(vec![ + ("value".to_string(), Value::Bytes(value)), + ("metadata".to_string(), Value::Bytes(metadata)), + ]); + + let mut reader = AvroArrowArrayReader::try_new([Ok(value)].into_iter(), &schema).unwrap(); + let batch = reader.next_batch(10).unwrap().unwrap(); + let variant_column = batch + .column_by_name("var") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(variant_column.column_by_name("metadata").is_some()); + assert_eq!( + VariantArray::try_new(batch.column_by_name("var").unwrap()) + .unwrap() + .value(0), + Variant::from("iceberg") + ); + } +} diff --git a/crates/core/src/avro_to_arrow/schema.rs b/crates/core/src/avro_to_arrow/schema.rs index 5b618b50f..4a8f2b1cb 100644 --- a/crates/core/src/avro_to_arrow/schema.rs +++ b/crates/core/src/avro_to_arrow/schema.rs @@ -15,15 +15,19 @@ // specific language governing permissions and limitations // under the License. -use crate::error::Result; +use crate::error::{CoreError, 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 parquet::variant::VariantType; use std::collections::HashMap; use std::sync::Arc; +const AVRO_LOGICAL_TYPE_KEY: &str = "logicalType"; +const HUDI_VARIANT_LOGICAL_TYPE: &str = "variant"; + /// Converts an avro schema to an arrow schema pub fn to_arrow_schema(avro_schema: &apache_avro::Schema) -> Result { let mut schema_fields = vec![]; @@ -125,7 +129,12 @@ fn schema_to_field_with_props( /*if let Some(aliases) = fields.aliases { props.insert("aliases", aliases); }*/ - schema_to_field_with_props(&field.schema, Some(&field.name), false, Some(props)) + schema_to_field_with_props( + &field.schema, + Some(&field.name), + field.is_nullable(), + Some(props), + ) }) .collect(); DataType::Struct(fields?) @@ -154,9 +163,91 @@ fn schema_to_field_with_props( let mut field = Field::new(name, field_type, nullable); field.set_metadata(props.unwrap_or_default()); + if is_hudi_variant_schema(schema) { + validate_hudi_variant_field(&field)?; + field.try_with_extension_type(VariantType)?; + } Ok(field) } +fn is_hudi_variant_schema(schema: &AvroSchema) -> bool { + match schema { + AvroSchema::Record(record_schema) => is_hudi_variant_record(record_schema), + AvroSchema::Union(union_schema) => { + let variants = union_schema.variants(); + variants.len() == 2 + && variants + .iter() + .any(|schema| matches!(schema, AvroSchema::Null)) + && variants.iter().any(is_hudi_variant_schema) + } + _ => false, + } +} + +fn is_hudi_variant_record(record_schema: &RecordSchema) -> bool { + matches!( + record_schema.attributes.get(AVRO_LOGICAL_TYPE_KEY), + Some(serde_json::Value::String(logical_type)) if logical_type == HUDI_VARIANT_LOGICAL_TYPE + ) +} + +fn validate_hudi_variant_field(field: &Field) -> Result<()> { + let DataType::Struct(fields) = field.data_type() else { + return Err(CoreError::Schema(format!( + "Hudi variant field '{}' must be an Avro record", + field.name() + ))); + }; + + let metadata = fields.iter().find(|child| child.name() == "metadata"); + match metadata { + Some(child) if is_variant_binary_field(child) => {} + Some(child) => { + return Err(CoreError::Schema(format!( + "Hudi variant field '{}' must have binary 'metadata', got {}", + field.name(), + child.data_type() + ))); + } + None => { + return Err(CoreError::Schema(format!( + "Hudi variant field '{}' must contain a 'metadata' field", + field.name() + ))); + } + } + + if let Some(value) = fields.iter().find(|child| child.name() == "value") + && !is_variant_binary_field(value) + { + return Err(CoreError::Schema(format!( + "Hudi variant field '{}' must have binary 'value', got {}", + field.name(), + value.data_type() + ))); + } + + if !fields + .iter() + .any(|child| matches!(child.name().as_str(), "value" | "typed_value")) + { + return Err(CoreError::Schema(format!( + "Hudi variant field '{}' must contain 'value' or 'typed_value'", + field.name() + ))); + } + + Ok(()) +} + +fn is_variant_binary_field(field: &Field) -> bool { + matches!( + field.data_type(), + DataType::Binary | DataType::LargeBinary | DataType::BinaryView + ) +} + fn default_field_name(dt: &DataType) -> &str { match dt { DataType::Null => "null", @@ -284,3 +375,100 @@ pub fn aliased(alias: &Alias, namespace: Option<&str>, default_namespace: Option } } } + +#[cfg(test)] +mod tests { + use super::*; + use parquet::variant::VariantType; + + fn root_schema(field_type: &str) -> AvroSchema { + AvroSchema::parse_str(&format!( + r#"{{ + "type": "record", + "name": "root", + "fields": [ + {{ + "name": "var", + "type": {field_type} + }} + ] + }}"# + )) + .unwrap() + } + + fn variant_record(fields: &str) -> String { + format!( + r#"{{ + "type": "record", + "name": "variant_record", + "logicalType": "variant", + "fields": [{fields}] + }}"# + ) + } + + #[test] + fn maps_hudi_variant_record_to_arrow_variant_extension() { + let schema = root_schema(&variant_record( + r#" + {"name": "metadata", "type": "bytes"}, + {"name": "value", "type": "bytes"} + "#, + )); + + let arrow_schema = to_arrow_schema(&schema).unwrap(); + let field = arrow_schema.field_with_name("var").unwrap(); + + field.try_extension_type::().unwrap(); + let DataType::Struct(fields) = field.data_type() else { + panic!("variant should be represented as an Arrow struct"); + }; + assert_eq!(fields[0].name(), "metadata"); + assert_eq!(fields[1].name(), "value"); + } + + #[test] + fn maps_nullable_hudi_variant_record_to_nullable_arrow_variant_extension() { + let schema = root_schema(&format!( + r#"[ "null", {} ]"#, + variant_record( + r#" + {"name": "metadata", "type": "bytes"}, + {"name": "value", "type": "bytes"} + "#, + ) + )); + + let arrow_schema = to_arrow_schema(&schema).unwrap(); + let field = arrow_schema.field_with_name("var").unwrap(); + + assert!(field.is_nullable()); + field.try_extension_type::().unwrap(); + } + + #[test] + fn rejects_hudi_variant_record_without_metadata() { + let schema = root_schema(&variant_record( + r#" + {"name": "value", "type": "bytes"} + "#, + )); + + let err = to_arrow_schema(&schema).unwrap_err(); + assert!(err.to_string().contains("must contain a 'metadata' field")); + } + + #[test] + fn rejects_hudi_variant_record_with_non_binary_value() { + let schema = root_schema(&variant_record( + r#" + {"name": "metadata", "type": "bytes"}, + {"name": "value", "type": "string"} + "#, + )); + + let err = to_arrow_schema(&schema).unwrap_err(); + assert!(err.to_string().contains("must have binary 'value'")); + } +} diff --git a/crates/core/src/file_group/base_file/parquet.rs b/crates/core/src/file_group/base_file/parquet.rs index 7c85833a9..29e738afd 100644 --- a/crates/core/src/file_group/base_file/parquet.rs +++ b/crates/core/src/file_group/base_file/parquet.rs @@ -189,8 +189,15 @@ impl BaseFileReader for ParquetBaseFileReader { #[cfg(test)] mod tests { use super::*; + use arrow_array::{ArrayRef, RecordBatch}; + use parquet::arrow::ArrowWriter; + use parquet::variant::{ + Variant, VariantArray, VariantArrayBuilder, VariantBuilderExt, VariantType, + }; + use std::fs::File; use std::fs::canonicalize; use std::path::Path; + use tempfile::TempDir; use url::Url; fn test_storage() -> Arc { @@ -199,6 +206,24 @@ mod tests { Storage::new_with_base_url(base_url).unwrap() } + fn temp_storage(temp_dir: &TempDir) -> Arc { + let base_url = Url::from_directory_path(temp_dir.path()).unwrap(); + Storage::new_with_base_url(base_url).unwrap() + } + + fn write_variant_parquet(path: &Path) { + let mut builder = VariantArrayBuilder::new(2); + builder.append_value("iceberg"); + builder.append_value("hudi"); + let array = builder.build(); + let schema = arrow_schema::Schema::new(vec![array.field("var")]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![ArrayRef::from(array)]).unwrap(); + let file = File::create(path).unwrap(); + let mut writer = ArrowWriter::try_new(file, batch.schema(), None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + } + #[tokio::test] async fn test_read_data_returns_all_rows() { let reader = ParquetBaseFileReader::new(test_storage()); @@ -280,4 +305,42 @@ mod tests { assert_eq!(meta.file_metadata().num_rows(), 5); assert!(!meta.row_groups().is_empty()); } + + #[tokio::test] + async fn test_read_variant_parquet_column() { + let temp_dir = tempfile::tempdir().unwrap(); + write_variant_parquet(&temp_dir.path().join("variant.parquet")); + let reader = ParquetBaseFileReader::new(temp_storage(&temp_dir)); + + let batch = reader + .read_data("variant.parquet", BaseFileReadOptions::default()) + .await + .unwrap(); + let schema = batch.schema(); + let field = schema.field_with_name("var").unwrap(); + field.try_extension_type::().unwrap(); + + let variant_array = VariantArray::try_new(batch.column_by_name("var").unwrap()).unwrap(); + assert_eq!(variant_array.value(0), Variant::from("iceberg")); + assert_eq!(variant_array.value(1), Variant::from("hudi")); + } + + #[tokio::test] + async fn test_project_variant_parquet_column() { + let temp_dir = tempfile::tempdir().unwrap(); + write_variant_parquet(&temp_dir.path().join("variant.parquet")); + let reader = ParquetBaseFileReader::new(temp_storage(&temp_dir)); + + let opts = BaseFileReadOptions::default().with_projection(["var"]); + let batch = reader.read_data("variant.parquet", opts).await.unwrap(); + + assert_eq!(batch.num_columns(), 1); + batch + .schema() + .field_with_name("var") + .unwrap() + .try_extension_type::() + .unwrap(); + assert!(VariantArray::try_new(batch.column_by_name("var").unwrap()).is_ok()); + } }