diff --git a/serde_arrow/src/internal/arrow/data_type.rs b/serde_arrow/src/internal/arrow/data_type.rs index a17405bd..580db657 100644 --- a/serde_arrow/src/internal/arrow/data_type.rs +++ b/serde_arrow/src/internal/arrow/data_type.rs @@ -12,6 +12,51 @@ pub struct Field { pub metadata: HashMap, } +impl PartialOrd for Field { + fn partial_cmp(&self, other: &Self) -> Option { + self.name.partial_cmp(&other.name) + } +} + +impl Field { + pub fn to_flattened_union_field(mut self, variant_name: &str) -> Self { + self.name = format!("{}::{}", variant_name, self.name); + self.nullable = true; + self + } + + fn from_flattened_union(&self) -> bool { + self.name.contains("::") + } + + pub fn union_variant_name(&self) -> Option<&str> { + if self.from_flattened_union() { + self.name.split("::").next() + } else { + None + } + } + + pub fn union_field_name(&self) -> Option { + if self.from_flattened_union() { + Some( + self.name + .split("::") + .skip(1) + .fold(String::new(), |acc: String, e| { + if acc.is_empty() { + String::from(e) + } else { + format!("{acc}::{e}") + } + }), + ) + } else { + None + } + } +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[non_exhaustive] pub enum DataType { diff --git a/serde_arrow/src/internal/schema/extensions/mod.rs b/serde_arrow/src/internal/schema/extensions/mod.rs index 9b11016e..426368b9 100644 --- a/serde_arrow/src/internal/schema/extensions/mod.rs +++ b/serde_arrow/src/internal/schema/extensions/mod.rs @@ -5,6 +5,7 @@ mod variable_shape_tensor_field; pub use bool8_field::Bool8Field; pub use fixed_shape_tensor_field::FixedShapeTensorField; +pub(crate) use utils::fix_dictionaries; pub use variable_shape_tensor_field::VariableShapeTensorField; const _: () = { diff --git a/serde_arrow/src/internal/schema/extensions/utils.rs b/serde_arrow/src/internal/schema/extensions/utils.rs index aae4d1c0..99e8df55 100644 --- a/serde_arrow/src/internal/schema/extensions/utils.rs +++ b/serde_arrow/src/internal/schema/extensions/utils.rs @@ -1,4 +1,7 @@ -use crate::internal::error::{fail, Result}; +use crate::internal::{ + arrow::{DataType, Field}, + error::{fail, Result}, +}; pub fn check_dim_names(ndim: usize, dim_names: &[String]) -> Result<()> { if dim_names.len() != ndim { @@ -56,3 +59,13 @@ impl std::fmt::Display for DebugRepr { write!(f, "{:?}", self.0) } } + +pub(crate) fn fix_dictionaries(field: &mut Field) { + if matches!(field.data_type, DataType::Dictionary(_, _, _)) { + field.nullable = true; + } else if let DataType::Struct(children) = &mut field.data_type { + for child in children { + fix_dictionaries(child); + } + } +} diff --git a/serde_arrow/src/internal/schema/from_samples/mod.rs b/serde_arrow/src/internal/schema/from_samples/mod.rs index ede6942c..50624eef 100644 --- a/serde_arrow/src/internal/schema/from_samples/mod.rs +++ b/serde_arrow/src/internal/schema/from_samples/mod.rs @@ -816,7 +816,10 @@ mod test { use serde::Serialize; use serde_json::{json, Value}; - use crate::internal::schema::{transmute_field, TracingOptions}; + use crate::internal::{ + schema::{transmute_field, TracingOptions}, + testing::{Coin, Number, Optionals, Payment}, + }; use super::*; @@ -920,4 +923,52 @@ mod test { expected, ); } + + #[test] + fn example_enum_as_struct_equal_to_struct_with_nullable_fields() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_samples(Number::sample_items().as_slice(), opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Number::expected_field()); + } + + #[test] + fn example_enum_as_struct_no_fields() { + // This should continue to maintain previously implemented behavior, serializing as a map + let opts = TracingOptions::default() + .enums_with_named_fields_as_structs(true) + .enums_without_data_as_strings(true); + + let enum_tracer = Tracer::from_samples(Coin::sample_items(), opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Coin::expected_field()); + } + + #[test] + #[should_panic] + fn example_enum_as_struct_no_fields_panics_when_opts_not_set() { + // This should continue to maintain previously implemented behavior, + // throwing an error because we detect Unions with no fields + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + + Tracer::from_samples(Coin::sample_items(), opts) + .unwrap() + .to_field() + .unwrap(); + } + + #[test] + fn example_enum_as_struct_all_fields_nullable() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_samples(Optionals::sample_items(), opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Optionals::expected_field()); + } + + #[test] + #[should_panic] + fn example_enum_as_struct_tuple_variants() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_samples(Payment::sample_items(), opts).unwrap(); + + // Currently panics when `to_schema()` is called on the variant tracer + enum_tracer.to_field().unwrap(); + } } diff --git a/serde_arrow/src/internal/schema/from_type/mod.rs b/serde_arrow/src/internal/schema/from_type/mod.rs index 38562591..7acdb77c 100644 --- a/serde_arrow/src/internal/schema/from_type/mod.rs +++ b/serde_arrow/src/internal/schema/from_type/mod.rs @@ -585,3 +585,59 @@ impl<'de, 'a> serde::de::Deserializer<'de> for IdentifierDeserializer<'a> { unimplemented!('de, deserialize_enum, _: &'static str, _: &'static [&'static str]); unimplemented!('de, deserialize_ignored_any); } + +#[cfg(test)] +mod test { + use crate::{ + internal::{ + schema::tracer::Tracer, + testing::{Coin, Number, Optionals, Payment}, + }, + schema::TracingOptions, + }; + + #[test] + fn example_enum_as_struct_equal_to_struct_with_nullable_fields() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_type::(opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Number::expected_field()); + } + + #[test] + fn example_enum_as_struct_no_fields() { + // This should continue to maintain previously implemented behavior, serializing as a map + let opts = TracingOptions::default() + .enums_with_named_fields_as_structs(true) + .enums_without_data_as_strings(true); + + let enum_tracer = Tracer::from_type::(opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Coin::expected_field()); + } + + #[test] + #[should_panic] + fn example_enum_as_struct_no_fields_panics_when_opts_not_set() { + // This should continue to maintain previously implemented behavior, + // throwing an error because we detect Unions with no fields + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + + Tracer::from_type::(opts).unwrap().to_field().unwrap(); + } + + #[test] + fn example_enum_as_struct_all_fields_nullable() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_type::(opts).unwrap(); + assert_eq!(enum_tracer.to_field().unwrap(), Optionals::expected_field()); + } + + #[test] + #[should_panic] + fn example_enum_as_struct_tuple_variants() { + let opts = TracingOptions::default().enums_with_named_fields_as_structs(true); + let enum_tracer = Tracer::from_type::(opts).unwrap(); + + // Currently panics when `to_schema()` is called on the variant tracer + enum_tracer.to_field().unwrap(); + } +} diff --git a/serde_arrow/src/internal/schema/mod.rs b/serde_arrow/src/internal/schema/mod.rs index 830b2ce4..5f6b29c6 100644 --- a/serde_arrow/src/internal/schema/mod.rs +++ b/serde_arrow/src/internal/schema/mod.rs @@ -482,7 +482,10 @@ fn validate_time64_field(field: &Field, unit: TimeUnit) -> Result<()> { fn validate_struct_field(field: &Field, children: &[Field]) -> Result<()> { // NOTE: do not check number of children: arrow-rs can 0 children, arrow2 not match get_strategy_from_metadata(&field.metadata)? { - None | Some(Strategy::MapAsStruct) | Some(Strategy::TupleAsStruct) => {} + None + | Some(Strategy::MapAsStruct) + | Some(Strategy::TupleAsStruct) + | Some(Strategy::EnumsWithNamedFieldsAsStructs) => {} Some(strategy) => fail!("invalid strategy for Struct field: {strategy}"), } for child in children { diff --git a/serde_arrow/src/internal/schema/strategy.rs b/serde_arrow/src/internal/schema/strategy.rs index 29162d64..1f743de1 100644 --- a/serde_arrow/src/internal/schema/strategy.rs +++ b/serde_arrow/src/internal/schema/strategy.rs @@ -63,6 +63,14 @@ pub enum Strategy { /// polars does not support them) /// MapAsStruct, + /// Serialize Rust enums that contain named field data as flattened structs + /// + /// This strategy is a workaround for the fact that Unions are not supported in parquet. + /// Currently, only Serialization is supported. Deserialization is not. + /// When writing out the enum, it will be flattened into a Struct with + /// a list of Fields, where the names of those Fields are the field prefixed with + /// the name of the variant. + EnumsWithNamedFieldsAsStructs, /// Mark a variant as unknown /// /// This strategy applies only to fields with DataType Null. If @@ -79,6 +87,7 @@ impl std::fmt::Display for Strategy { Self::NaiveStrAsDate64 => write!(f, "NaiveStrAsDate64"), Self::TupleAsStruct => write!(f, "TupleAsStruct"), Self::MapAsStruct => write!(f, "MapAsStruct"), + Self::EnumsWithNamedFieldsAsStructs => write!(f, "EnumsWithNamedFieldsAsStructs"), Self::UnknownVariant => write!(f, "UnknownVariant"), } } @@ -108,6 +117,7 @@ impl FromStr for Strategy { "NaiveStrAsDate64" => Ok(Self::NaiveStrAsDate64), "TupleAsStruct" => Ok(Self::TupleAsStruct), "MapAsStruct" => Ok(Self::MapAsStruct), + "EnumsWithNamedFieldsAsStructs" => Ok(Self::EnumsWithNamedFieldsAsStructs), "UnknownVariant" => Ok(Self::UnknownVariant), _ => fail!("Unknown strategy {s}"), } diff --git a/serde_arrow/src/internal/schema/tracer.rs b/serde_arrow/src/internal/schema/tracer.rs index daa460af..c45a77a0 100644 --- a/serde_arrow/src/internal/schema/tracer.rs +++ b/serde_arrow/src/internal/schema/tracer.rs @@ -7,8 +7,8 @@ use crate::internal::{ arrow::{DataType, Field, UnionMode}, error::{fail, set_default, Context, Result}, schema::{ - DataTypeDisplay, Overwrites, SerdeArrowSchema, Strategy, TracingMode, TracingOptions, - STRATEGY_KEY, + extensions::fix_dictionaries, DataTypeDisplay, Overwrites, SerdeArrowSchema, Strategy, + TracingMode, TracingOptions, STRATEGY_KEY, }, }; @@ -101,32 +101,52 @@ impl Tracer { Self::Unknown(UnknownTracer::new(name, path, options)) } + fn schema_tracing_error( + failed_data_type: impl std::fmt::Display, + tracing_mode: TracingMode, + ) -> Result { + fail!( + concat!( + "Schema tracing is not directly supported for the root data type {failed_data_type}. ", + "Only struct-like types are supported as root types in schema tracing. ", + "{mitigation}", + ), + failed_data_type = failed_data_type, + mitigation = match tracing_mode { + TracingMode::FromType => { + "Consider using the `Item` wrapper, i.e., `::from_type>()`." + } + TracingMode::FromSamples => { + "Consider using the `Items` wrapper, i.e., `::from_samples(Items(samples))`." + } + TracingMode::Unknown => "Consider using the `Item` / `Items` wrappers.", + }, + + ) + } + /// Convert the traced schema into a schema object pub fn to_schema(&self) -> Result { let root = self.to_field()?; if root.nullable { - fail!("The root type cannot be nullable"); + fail!("The root type cannot be nullable: {root:#?}"); } let tracing_mode = dispatch_tracer!(self, tracer => tracer.options.tracing_mode); let fields = match root.data_type { - DataType::Struct(children) => children, + DataType::Struct(children) => { + if let Some(strategy) = root.metadata.get(STRATEGY_KEY) { + if *strategy == Strategy::EnumsWithNamedFieldsAsStructs.to_string() { + return Self::schema_tracing_error("Union", tracing_mode); + } + } + + children + } DataType::Null => fail!("No records found to determine schema"), - dt => fail!( - concat!( - "Schema tracing is not directly supported for the root data type {dt}. ", - "Only struct-like types are supported as root types in schema tracing. ", - "{mitigation}", - ), - dt = DataTypeDisplay(&dt), - mitigation = match tracing_mode { - TracingMode::FromType => "Consider using the `Item` wrapper, i.e., `::from_type>()`.", - TracingMode::FromSamples => "Consider using the `Items` wrapper, i.e., `::from_samples(Items(samples))`.", - TracingMode::Unknown => "Consider using the `Item` / `Items` wrappers.", - }, - ), + dt => return Self::schema_tracing_error(DataTypeDisplay(&dt), tracing_mode), }; Ok(SerdeArrowSchema { fields }) @@ -1065,20 +1085,53 @@ impl UnionTracer { } } - let mut fields = Vec::new(); - for (idx, variant) in self.variants.iter().enumerate() { - if let Some(variant) = variant { - fields.push((i8::try_from(idx)?, variant.tracer.to_field()?)); - } else { - fields.push((i8::try_from(idx)?, unknown_variant_field())); - }; + let data_type: DataType; + let mut metadata = HashMap::new(); + + if self.options.enums_with_named_fields_as_structs { + metadata.insert( + STRATEGY_KEY.to_string(), + Strategy::EnumsWithNamedFieldsAsStructs.to_string(), + ); + let mut fields = BTreeMap::new(); + + // For this option, we want to merge the variant children up one level, combining the names + // For each variant with name variant_name + // For each variant_field with field_name + // Add field {variant_name}::{field_name} -> variant_field.to_field() that is nullable + + for variant in &self.variants { + if let Some(variant) = variant { + let schema = variant.tracer.to_schema()?; + for field in schema.fields { + let mut flat_field = field.to_flattened_union_field(variant.name.as_str()); + fix_dictionaries(&mut flat_field); + fields.insert(flat_field.name.to_string(), flat_field); + } + } else { + let uf = unknown_variant_field(); + fields.insert(uf.name, unknown_variant_field()); + }; + } + + data_type = DataType::Struct(fields.into_values().collect()); + } else { + let mut fields = Vec::new(); + for (idx, variant) in self.variants.iter().enumerate() { + if let Some(variant) = variant { + fields.push((i8::try_from(idx)?, variant.tracer.to_field()?)); + } else { + fields.push((i8::try_from(idx)?, unknown_variant_field())); + }; + } + data_type = DataType::Union(fields, UnionMode::Dense); } Ok(Field { name: self.name.to_owned(), - data_type: DataType::Union(fields, UnionMode::Dense), + data_type, nullable: self.nullable, - metadata: HashMap::new(), + metadata, }) } diff --git a/serde_arrow/src/internal/schema/tracing_options.rs b/serde_arrow/src/internal/schema/tracing_options.rs index 86d09474..276ab84b 100644 --- a/serde_arrow/src/internal/schema/tracing_options.rs +++ b/serde_arrow/src/internal/schema/tracing_options.rs @@ -220,6 +220,28 @@ pub struct TracingOptions { /// Internal field to improve error messages for the different tracing /// functions pub(crate) tracing_mode: TracingMode, + + /// Whether to encode enums with data as structs + /// + /// If `false` enums with data are encoded as Union arrays. + /// If `true` enums with data are encoded as Structs. + /// + /// ``` + /// # fn main() -> serde_arrow::Result<()> { + /// # use serde_arrow::_impl::arrow; + /// # use arrow::datatypes::{FieldRef, Field, DataType, TimeUnit}; + /// # use serde_arrow::schema::{SchemaLike, TracingOptions}; + /// # use serde::{Serialize, Deserialize}; + /// #[derive(Serialize, Deserialize)] + /// enum Number { + /// Real { value: f32 }, + /// Complex { i: f32, j: f32 }, + /// } + /// let options = TracingOptions::default().enums_with_named_fields_as_structs(true); + /// let fields = Tracer::from_type::(options)?; + /// # } + /// ``` + pub enums_with_named_fields_as_structs: bool, } impl Default for TracingOptions { @@ -232,6 +254,7 @@ impl Default for TracingOptions { guess_dates: false, from_type_budget: 100, enums_without_data_as_strings: false, + enums_with_named_fields_as_structs: false, overwrites: Overwrites::default(), sequence_as_large_list: true, string_as_large_utf8: true, @@ -299,6 +322,12 @@ impl TracingOptions { self } + /// Set [`enums_with_named_fields_as_structs`](#structfield.enums_with_named_fields_as_structs) + pub fn enums_with_named_fields_as_structs(mut self, value: bool) -> Self { + self.enums_with_named_fields_as_structs = value; + self + } + /// Add an overwrite to [`overwrites`](#structfield.overwrites) pub fn overwrite, F: Serialize>(mut self, path: P, field: F) -> Result { self.overwrites.0.insert( diff --git a/serde_arrow/src/internal/serialization/array_builder.rs b/serde_arrow/src/internal/serialization/array_builder.rs index 75eb6ef5..7d3b2b68 100644 --- a/serde_arrow/src/internal/serialization/array_builder.rs +++ b/serde_arrow/src/internal/serialization/array_builder.rs @@ -13,10 +13,10 @@ use super::{ date64_builder::Date64Builder, decimal_builder::DecimalBuilder, dictionary_utf8_builder::DictionaryUtf8Builder, duration_builder::DurationBuilder, fixed_size_binary_builder::FixedSizeBinaryBuilder, - fixed_size_list_builder::FixedSizeListBuilder, float_builder::FloatBuilder, - int_builder::IntBuilder, list_builder::ListBuilder, map_builder::MapBuilder, - null_builder::NullBuilder, simple_serializer::SimpleSerializer, struct_builder::StructBuilder, - time_builder::TimeBuilder, union_builder::UnionBuilder, + fixed_size_list_builder::FixedSizeListBuilder, flattened_union_builder::FlattenedUnionBuilder, + float_builder::FloatBuilder, int_builder::IntBuilder, list_builder::ListBuilder, + map_builder::MapBuilder, null_builder::NullBuilder, simple_serializer::SimpleSerializer, + struct_builder::StructBuilder, time_builder::TimeBuilder, union_builder::UnionBuilder, unknown_variant_builder::UnknownVariantBuilder, utf8_builder::Utf8Builder, }; @@ -53,6 +53,7 @@ pub enum ArrayBuilder { LargeUtf8(Utf8Builder), DictionaryUtf8(DictionaryUtf8Builder), Union(UnionBuilder), + FlattenedUnion(FlattenedUnionBuilder), UnknownVariant(UnknownVariantBuilder), } @@ -90,6 +91,7 @@ macro_rules! dispatch { $wrapper::Struct($name) => $expr, $wrapper::DictionaryUtf8($name) => $expr, $wrapper::Union($name) => $expr, + $wrapper::FlattenedUnion($name) => $expr, $wrapper::UnknownVariant($name) => $expr, } }; diff --git a/serde_arrow/src/internal/serialization/flattened_union_builder.rs b/serde_arrow/src/internal/serialization/flattened_union_builder.rs new file mode 100644 index 00000000..3f647225 --- /dev/null +++ b/serde_arrow/src/internal/serialization/flattened_union_builder.rs @@ -0,0 +1,119 @@ +use std::collections::BTreeMap; + +use crate::internal::{ + arrow::{Array, FieldMeta, StructArray}, + error::{fail, set_default, try_, Context, ContextSupport, Result}, +}; + +use super::{array_builder::ArrayBuilder, simple_serializer::SimpleSerializer}; + +#[derive(Debug, Clone)] +pub struct FlattenedUnionBuilder { + path: String, + fields: Vec<(ArrayBuilder, FieldMeta)>, + row_count: usize, +} + +impl FlattenedUnionBuilder { + pub fn new(path: String, fields: Vec<(ArrayBuilder, FieldMeta)>) -> Self { + Self { + path, + fields, + row_count: 0, + } + } + + pub fn take(&mut self) -> ArrayBuilder { + ArrayBuilder::FlattenedUnion(Self { + path: self.path.clone(), + fields: self + .fields + .iter_mut() + .map(|(field, meta)| (field.take(), meta.clone())) + .collect(), + row_count: self.row_count, + }) + } + + pub fn is_nullable(&self) -> bool { + false + } + + pub fn into_array(self) -> Result { + let mut fields = BTreeMap::new(); + + for (builder, meta) in self.fields.into_iter() { + let ArrayBuilder::Struct(builder) = builder else { + fail!("Attempting to flatten a not-struct builder: {builder:?}"); + }; + + for (sub_builder, mut sub_meta) in builder.fields.into_iter() { + // TODO: this mirrors the field name structure in the tracer but represents + // implementation details crossing boundaries. Is there another way? + // Name change is currently needed for struct field lookup to work correctly. + + sub_meta.name = format!("{}::{}", meta.name, sub_meta.name); + fields.insert( + sub_meta.name.to_owned(), + (sub_builder.into_array()?, sub_meta), + ); + } + } + + Ok(Array::Struct(StructArray { + len: self.row_count, + fields: fields.into_values().collect(), + // assuming this is OK to hardcode because empirically, + // the validity of struct with nullable fields was always None + validity: None, + })) + } +} + +impl FlattenedUnionBuilder { + pub fn serialize_variant(&mut self, variant_index: u32) -> Result<&mut ArrayBuilder> { + let variant_index = variant_index as usize; + + // don't serialize any variant not selected + for (idx, (builder, _meta)) in self.fields.iter_mut().enumerate() { + if idx != variant_index { + builder.serialize_none()?; + } + } + + let Some((variant_builder, _variant_meta)) = self.fields.get_mut(variant_index) else { + fail!("Could not find variant {variant_index} in Union"); + }; + + self.row_count += 1; + + Ok(variant_builder) + } +} + +impl Context for FlattenedUnionBuilder { + fn annotate(&self, annotations: &mut BTreeMap) { + set_default(annotations, "field", &self.path); + set_default(annotations, "data_type", "Struct(..)"); + } +} + +impl SimpleSerializer for FlattenedUnionBuilder { + fn serialize_struct_variant_start<'this>( + &'this mut self, + _: &'static str, + variant_index: u32, + variant: &'static str, + len: usize, + ) -> Result<&'this mut ArrayBuilder> { + let mut ctx = BTreeMap::new(); + self.annotate(&mut ctx); + + try_(|| { + let variant_builder = self.serialize_variant(variant_index)?; + variant_builder.serialize_struct_start(variant, len)?; + Ok(variant_builder) + }) + .ctx(&ctx) + } +} diff --git a/serde_arrow/src/internal/serialization/mod.rs b/serde_arrow/src/internal/serialization/mod.rs index f6af48eb..8198e634 100644 --- a/serde_arrow/src/internal/serialization/mod.rs +++ b/serde_arrow/src/internal/serialization/mod.rs @@ -10,6 +10,7 @@ pub mod dictionary_utf8_builder; pub mod duration_builder; pub mod fixed_size_binary_builder; pub mod fixed_size_list_builder; +pub mod flattened_union_builder; pub mod float_builder; pub mod int_builder; pub mod list_builder; diff --git a/serde_arrow/src/internal/serialization/outer_sequence_builder.rs b/serde_arrow/src/internal/serialization/outer_sequence_builder.rs index 58bbf1ba..7d27dcf3 100644 --- a/serde_arrow/src/internal/serialization/outer_sequence_builder.rs +++ b/serde_arrow/src/internal/serialization/outer_sequence_builder.rs @@ -3,13 +3,14 @@ use std::collections::{BTreeMap, HashMap}; use serde::Serialize; use crate::internal::{ - arrow::{DataType, Field, TimeUnit}, + arrow::{DataType, Field, FieldMeta, TimeUnit}, error::{fail, Context, ContextSupport, Result}, schema::{get_strategy_from_metadata, SerdeArrowSchema, Strategy}, serialization::{ binary_builder::BinaryBuilder, duration_builder::DurationBuilder, fixed_size_binary_builder::FixedSizeBinaryBuilder, fixed_size_list_builder::FixedSizeListBuilder, + flattened_union_builder::FlattenedUnionBuilder, }, utils::{btree_map, meta_from_field, ChildName, Mut}, }; @@ -121,7 +122,7 @@ fn build_struct(path: String, struct_fields: &[Field], nullable: bool) -> Result StructBuilder::new(path, fields, nullable) } -fn build_builder(path: String, field: &Field) -> Result { +pub(crate) fn build_builder(path: String, field: &Field) -> Result { use {ArrayBuilder as A, DataType as T}; let ctx: BTreeMap = btree_map!("field" => path.clone()); @@ -226,7 +227,50 @@ fn build_builder(path: String, field: &Field) -> Result { .ctx(&ctx)?, ) } - T::Struct(children) => A::Struct(build_struct(path, children, field.nullable)?), + T::Struct(children) => { + if let Some(Strategy::EnumsWithNamedFieldsAsStructs) = + get_strategy_from_metadata(&field.metadata)? + { + let mut related_fields: BTreeMap<&str, Vec> = BTreeMap::new(); + let mut builders: Vec<(ArrayBuilder, FieldMeta)> = Vec::new(); + + for field in children { + let Some(variant_name) = field.union_variant_name() else { + todo!("union variant did not have a name"); + }; + + let Some(field_name) = field.union_field_name() else { + todo!("union field did not have a name"); + }; + + let mut new_field = field.clone(); + new_field.name = field_name; + + related_fields + .entry(variant_name) + .or_default() + .push(new_field); + } + + for (variant_name, fields) in related_fields { + let builder = build_struct( + format!("{}.{}", path.to_owned(), variant_name), + fields.as_slice(), + true, + )? + .take(); + + let mut meta = meta_from_field(field.clone()); + meta.name = variant_name.to_owned(); + + builders.push((builder, meta)); + } + + A::FlattenedUnion(FlattenedUnionBuilder::new(path, builders)) + } else { + A::Struct(build_struct(path, children, field.nullable)?) + } + } T::Dictionary(key, value, _) => { let key_path = format!("{path}.key"); let key_field = Field { diff --git a/serde_arrow/src/internal/serialization/simple_serializer.rs b/serde_arrow/src/internal/serialization/simple_serializer.rs index 80183e0b..2de770b2 100644 --- a/serde_arrow/src/internal/serialization/simple_serializer.rs +++ b/serde_arrow/src/internal/serialization/simple_serializer.rs @@ -169,7 +169,7 @@ pub trait SimpleSerializer: Sized + Context { fn serialize_struct_start(&mut self, name: &'static str, len: usize) -> Result<()> { fail!( in self, - "serialize_start_start is not supported", + "serialize_struct_start is not supported", ) } diff --git a/serde_arrow/src/internal/testing.rs b/serde_arrow/src/internal/testing.rs index 1b6ac8b0..9fc20e82 100644 --- a/serde_arrow/src/internal/testing.rs +++ b/serde_arrow/src/internal/testing.rs @@ -1,10 +1,14 @@ //! Support for tests use core::str; +use std::collections::HashMap; use crate::internal::{ - arrow::{Array, BytesArray}, + arrow::{Array, BytesArray, DataType, Field}, error::{fail, Error, Result}, }; +use crate::schema::{Strategy, STRATEGY_KEY}; + +use serde::{Deserialize, Serialize}; pub fn assert_error_contains(actual: &Result, expected: &str) { let Err(actual) = actual else { @@ -75,3 +79,173 @@ where Ok(Some(str::from_utf8(data)?)) } + +fn enum_with_named_fields_metadata() -> HashMap { + HashMap::from([( + STRATEGY_KEY.to_string(), + Strategy::EnumsWithNamedFieldsAsStructs.to_string(), + )]) +} + +// Simple enum test structure for schema from_type/from_samples unit testing +#[derive(Serialize, Deserialize)] +pub(crate) enum Number { + Real { value: f32 }, + Complex { i: f32, j: f32 }, +} + +impl Number { + pub(crate) fn sample_items() -> Vec { + vec![ + Number::Real { value: 1.0 }, + Number::Complex { i: 0.5, j: 0.5 }, + ] + } + + pub(crate) fn expected_field() -> Field { + Field { + name: "$".to_string(), + data_type: DataType::Struct(vec![ + Field { + name: "Complex::i".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Complex::j".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Real::value".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + ]), + nullable: false, + metadata: enum_with_named_fields_metadata(), + } + } +} + +// No data test enum +#[derive(Serialize, Deserialize)] +pub(crate) enum Coin { + Heads, + Tails, +} + +impl Coin { + pub(crate) fn sample_items() -> Vec { + vec![Coin::Heads, Coin::Tails] + } + + pub(crate) fn expected_field() -> Field { + Field { + name: "$".to_string(), + data_type: DataType::Dictionary( + Box::new(DataType::UInt32), + Box::new(DataType::LargeUtf8), + false, + ), + nullable: false, + metadata: HashMap::new(), + } + } +} + +// Optional variant field test enum +#[derive(Serialize, Deserialize)] +pub(crate) enum Optionals { + Something { + more: Option, + less: Option, + }, + Else { + one: Option, + another: Option, + }, +} + +impl Optionals { + pub(crate) fn sample_items() -> Vec { + vec![ + Optionals::Something { + more: Some(1), + less: None, + }, + Optionals::Something { + more: None, + less: Some(0), + }, + Optionals::Else { + one: None, + another: Some(0), + }, + Optionals::Else { + one: Some(1), + another: None, + }, + ] + } + + pub(crate) fn expected_field() -> Field { + Field { + name: "$".to_string(), + data_type: DataType::Struct(vec![ + Field { + name: "Else::another".to_string(), + data_type: DataType::UInt64, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Else::one".to_string(), + data_type: DataType::UInt64, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Something::less".to_string(), + data_type: DataType::UInt64, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Something::more".to_string(), + data_type: DataType::UInt64, + nullable: true, + metadata: HashMap::new(), + }, + ]), + nullable: false, + metadata: enum_with_named_fields_metadata(), + } + } +} + +// Tuple variant test enum +#[derive(Serialize, Deserialize)] +pub(crate) enum Payment { + Cash(f32), // amount + Check(String, f32), // name, amount + CreditCard(String, f32, [u8; 16], String), // name, amount, cc number, exp +} + +impl Payment { + pub(crate) fn sample_items() -> Vec { + vec![ + Payment::Cash(0.42), + Payment::Check("Bob".to_string(), 0.42), + Payment::CreditCard( + "Sue".to_string(), + 0.42, + [1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6], + "01/2024".to_string(), + ), + ] + } +} diff --git a/serde_arrow/src/test_with_arrow/impls/flattened_union.rs b/serde_arrow/src/test_with_arrow/impls/flattened_union.rs new file mode 100644 index 00000000..2759a1d0 --- /dev/null +++ b/serde_arrow/src/test_with_arrow/impls/flattened_union.rs @@ -0,0 +1,371 @@ +use std::collections::HashMap; + +use crate::{ + internal::{ + array_builder::ArrayBuilder, + arrow::{Array, DataType, Field}, + schema::{SchemaLike, TracingOptions}, + }, + schema::SerdeArrowSchema, + Serializer, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize)] +struct Number { + v: Value, +} + +#[derive(Serialize, Deserialize)] +enum Value { + Real { value: f32 }, + Complex { i: f32, j: f32 }, + Whole { value: usize }, +} + +fn number_field() -> Field { + Field { + name: "v".to_string(), + data_type: DataType::Struct(vec![ + Field { + name: "Complex::i".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Complex::j".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Real::value".to_string(), + data_type: DataType::Float32, + nullable: true, + metadata: HashMap::new(), + }, + Field { + name: "Whole::value".to_string(), + data_type: DataType::UInt64, + nullable: true, + metadata: HashMap::new(), + }, + ]), + nullable: false, + metadata: HashMap::from([( + "SERDE_ARROW:strategy".to_string(), + "EnumsWithNamedFieldsAsStructs".to_string(), + )]), + } +} + +fn number_schema() -> SerdeArrowSchema { + let options = TracingOptions::default() + .allow_null_fields(true) + .enums_with_named_fields_as_structs(true); + + SerdeArrowSchema::from_type::(options).unwrap() +} + +fn number_data() -> Vec { + vec![ + Number { + v: Value::Real { value: 0.0 }, + }, + Number { + v: Value::Complex { i: 0.5, j: 0.5 }, + }, + Number { + v: Value::Whole { value: 5 }, + }, + ] +} + +#[test] +fn test_build_flattened_union_builder() { + let mut builder = ArrayBuilder::new(number_schema()).unwrap(); + + // One struct in the array + let arrays = builder.build_arrays().unwrap(); + + assert_eq!(arrays.len(), 1); + + let array = &arrays[0]; + + let Array::Struct(ref struct_array) = array else { + panic!("expected a struct array, found {array:#?}"); + }; + + // Should be a single struct array with 4 fields: Complex::i, Complex::j, Real::value, Whole::value + assert_eq!( + struct_array.fields.len(), + 4, + "contained {} fields", + struct_array.fields.len() + ); + + let (first_field, meta) = &struct_array.fields[0]; + assert_eq!(meta.name, "Complex::i"); + assert!(matches!(first_field, Array::Float32(_))); + + let (second_field, meta) = &struct_array.fields[1]; + assert_eq!(meta.name, "Complex::j"); + assert!(matches!(second_field, Array::Float32(_))); + + let (third_field, meta) = &struct_array.fields[2]; + assert_eq!(meta.name, "Real::value"); + assert!(matches!(third_field, Array::Float32(_))); + + let (fourth_field, meta) = &struct_array.fields[3]; + assert_eq!(meta.name, "Whole::value"); + assert!(matches!(fourth_field, Array::UInt64(_))); +} + +#[test] +fn test_serialize_flattened_union_builder() { + let field = number_field(); + let data = number_data(); + let schema = SerdeArrowSchema { + fields: vec![field], + }; + + let api_builder = ArrayBuilder::new(schema).expect("failed to create api array builder"); + let serializer = Serializer::new(api_builder); + data.serialize(serializer) + .expect("failed to serialize") + .into_inner() + .to_arrow() + .expect("failed to serialize to arrow"); +} + +#[test] +fn test_record_batch_flattened_union_builder() { + let field = number_field(); + let data = number_data(); + let schema = SerdeArrowSchema { + fields: vec![field], + }; + + let api_builder = ArrayBuilder::new(schema).expect("failed to create api array builder"); + let serializer = Serializer::new(api_builder); + data.serialize(serializer) + .expect("failed to serialize") + .into_inner() + .to_record_batch() + .expect("failed to create record batch"); +} + +#[derive(Serialize, Deserialize)] +struct ComplexMessage { + data: MsgData, +} + +#[derive(Serialize, Deserialize)] +enum MsgData { + One { data: usize }, + Two { opts: MsgOptions }, +} + +#[derive(Serialize, Deserialize)] +struct MsgOptions { + loc: Location, +} + +#[derive(Serialize, Deserialize, Default)] +enum Location { + #[default] + Left, + Right, +} + +fn nested_enum_schema() -> SerdeArrowSchema { + let options = TracingOptions::default() + .allow_null_fields(true) + .enums_without_data_as_strings(true) + .enums_with_named_fields_as_structs(true); + + SerdeArrowSchema::from_type::(options).unwrap() +} + +fn nested_enum_data() -> Vec { + vec![ + ComplexMessage { + data: MsgData::One { data: 3 }, + }, + ComplexMessage { + data: MsgData::Two { + opts: MsgOptions { + loc: Location::Right, + }, + }, + }, + ] +} + +#[test] +fn test_flattened_union_with_nested_enum() { + let mut builder = ArrayBuilder::new(nested_enum_schema()).unwrap(); + + // One struct in the array + let arrays = builder.build_arrays().unwrap(); + + println!("{arrays:#?}"); + + assert_eq!(arrays.len(), 1); + + let array = &arrays[0]; + + let Array::Struct(ref _struct_array) = array else { + panic!("expected a struct array, found {array:#?}"); + }; + + let serializer = Serializer::new(builder); + + let result = nested_enum_data() + .serialize(serializer) + .expect("failed to serialize") + .into_inner() + .to_arrow() + .expect("failed to serialize to arrow"); + + println!("arrow: {result:#?}"); +} + +#[derive(Serialize, Deserialize)] +struct OuterSkipped { + meta: u32, + data: InnerSkipped, +} + +#[derive(Serialize, Deserialize)] +enum InnerSkipped { + One { + data: usize, + }, + Two { + vector: [i32; 3], + #[serde(skip)] + location: Location, + }, + Three { + // x: Option, + // y: Option, + y: InnerData, + }, +} + +#[derive(Serialize, Deserialize, Default)] +struct InnerData { + field1: usize, + field2: i32, + field3: u32, +} + +fn skipped_schema() -> SerdeArrowSchema { + let options = TracingOptions::default() + .allow_null_fields(true) + .enums_without_data_as_strings(true) + .enums_with_named_fields_as_structs(true); + + SerdeArrowSchema::from_type::(options).unwrap() +} + +fn skipped_data() -> Vec { + vec![ + OuterSkipped { + meta: 1, + data: InnerSkipped::One { data: 1 }, + }, + OuterSkipped { + meta: 2, + data: InnerSkipped::Two { + vector: [2, 2, 2], + location: Location::Right, + }, + }, + // OuterSkipped { + // meta: 3, + // data: InnerSkipped::Three { + // /*x: None,*/ y: None, + // }, + // }, + OuterSkipped { + meta: 4, + data: InnerSkipped::Three { + // x: None, + // y: Some(InnerData { + // field1: 99, + // field2: -99, + // field3: 99, + // }), + y: InnerData { + field1: 99, + field2: -99, + field3: 99, + }, + }, + }, + ] +} + +#[test] +fn test_flattened_union_with_nested_skipped_fields() { + let mut builder = ArrayBuilder::new(skipped_schema()).unwrap(); + + // One struct in the array + let arrays = builder.build_arrays().unwrap(); + + println!("{arrays:#?}"); + + assert_eq!(arrays.len(), 2); + + let Array::UInt32(_) = &arrays[0] else { + panic!("expected a int array, found {:#?}", &arrays[0]); + }; + + let Array::Struct(_) = &arrays[1] else { + panic!("expected a struct array, found {:#?}", &arrays[1]); + }; + + let serializer = Serializer::new(builder); + + let result = skipped_data() + .serialize(serializer) + .expect("failed to serialize") + .into_inner() + .to_arrow() + .expect("failed to serialize to arrow"); + + println!("arrow: {result:#?}"); + + // TODO: I think this repros the None issue! + + /* + -- child 1: "Three::y" (Struct([Field { name: "field1", data_type: UInt64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "field2", data_type: Int32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, Field { name: "field3", data_type: UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }])) + StructArray + [ + -- child 0: "field1" (UInt64) + PrimitiveArray + [ + 0, + 0, + 0, + ] + -- child 1: "field2" (Int32) + PrimitiveArray + [ + 0, + 0, + 0, + ] + -- child 2: "field3" (UInt32) + PrimitiveArray + [ + 0, + 0, + 0, + ] + ] + */ +} diff --git a/serde_arrow/src/test_with_arrow/impls/mod.rs b/serde_arrow/src/test_with_arrow/impls/mod.rs index 7bbef5d5..dcf02fe4 100644 --- a/serde_arrow/src/test_with_arrow/impls/mod.rs +++ b/serde_arrow/src/test_with_arrow/impls/mod.rs @@ -6,6 +6,7 @@ mod chrono; mod dictionary; mod examples; mod fixed_size_list; +mod flattened_union; mod jiff; mod json_values; mod list;