diff --git a/arrow-integration-test/src/field.rs b/arrow-integration-test/src/field.rs index 68019cc5377d..f0b01bd46f4b 100644 --- a/arrow-integration-test/src/field.rs +++ b/arrow-integration-test/src/field.rs @@ -21,7 +21,19 @@ use arrow::error::{ArrowError, Result}; use std::collections::HashMap; use std::sync::Arc; +/// The metadata key of an extension type name +const EXTENSION_TYPE_NAME_KEY: &str = "ARROW:extension:name"; + /// Parse a `Field` definition from a JSON representation. +/// +/// The optional `metadata` of a field's `dictionary` holds the metadata of the +/// dictionary value type, for example its extension type. [`DataType::Dictionary`] +/// cannot carry metadata for its value type (see +/// ), so it is mapped to the +/// metadata of the returned [`Field`] in the same way as the arrow-ipc readers +/// map `DictionaryEncoding.custom_metadata`: it is copied into the field's +/// metadata, without replacing keys already present there, unless the field +/// has its own `ARROW:extension:name`, in which case it is dropped. pub fn field_from_json(json: &serde_json::Value) -> Result { use serde_json::Value; match *json { @@ -49,64 +61,9 @@ pub fn field_from_json(json: &serde_json::Value) -> Result { }; // Referenced example file: testing/data/arrow-ipc-stream/integration/1.0.0-littleendian/generated_custom_metadata.json.gz - let metadata = match map.get("metadata") { - Some(Value::Array(values)) => { - let mut res: HashMap = HashMap::default(); - for value in values { - match value.as_object() { - Some(map) => { - if map.len() != 2 { - return Err(ArrowError::ParseError( - "Field 'metadata' must have exact two entries for each key-value map".to_string(), - )); - } - if let (Some(k), Some(v)) = (map.get("key"), map.get("value")) { - if let (Some(k_str), Some(v_str)) = (k.as_str(), v.as_str()) { - res.insert( - k_str.to_string().clone(), - v_str.to_string().clone(), - ); - } else { - return Err(ArrowError::ParseError( - "Field 'metadata' must have map value of string type" - .to_string(), - )); - } - } else { - return Err(ArrowError::ParseError("Field 'metadata' lacks map keys named \"key\" or \"value\"".to_string())); - } - } - _ => { - return Err(ArrowError::ParseError( - "Field 'metadata' contains non-object key-value pair" - .to_string(), - )); - } - } - } - res - } - // We also support map format, because Schema's metadata supports this. - // See https://github.com/apache/arrow/pull/5907 - Some(Value::Object(values)) => { - let mut res: HashMap = HashMap::default(); - for (k, v) in values { - if let Some(str_value) = v.as_str() { - res.insert(k.clone(), str_value.to_string().clone()); - } else { - return Err(ArrowError::ParseError(format!( - "Field 'metadata' contains non-string value for key {k}" - ))); - } - } - res - } - Some(_) => { - return Err(ArrowError::ParseError( - "Field `metadata` is not json array".to_string(), - )); - } - _ => HashMap::default(), + let mut metadata = match map.get("metadata") { + Some(metadata) => metadata_from_json(metadata, "metadata")?, + None => HashMap::default(), }; // if data_type is a struct or list, get its children @@ -278,6 +235,16 @@ pub fn field_from_json(json: &serde_json::Value) -> Result { )); } }; + // metadata for the dictionary value type, see above + if let Some(value_metadata) = dictionary.get("metadata") { + let value_metadata = + metadata_from_json(value_metadata, "dictionary.metadata")?; + if !metadata.contains_key(EXTENSION_TYPE_NAME_KEY) { + for (k, v) in value_metadata { + metadata.entry(k).or_insert(v); + } + } + } DataType::Dictionary(Box::new(index_type), Box::new(data_type)) } _ => data_type, @@ -294,6 +261,69 @@ pub fn field_from_json(json: &serde_json::Value) -> Result { } } +/// Parse the key-value metadata in the field attribute `attribute`. +/// +/// The metadata is a list of `{"key": ..., "value": ...}` objects. A map from +/// keys to values is also accepted, because Schema's metadata supports this. +/// See +fn metadata_from_json( + json: &serde_json::Value, + attribute: &str, +) -> Result> { + use serde_json::Value; + match json { + Value::Array(values) => { + let mut res: HashMap = HashMap::default(); + for value in values { + match value.as_object() { + Some(map) => { + if map.len() != 2 { + return Err(ArrowError::ParseError(format!( + "Field '{attribute}' must have exact two entries for each key-value map" + ))); + } + if let (Some(k), Some(v)) = (map.get("key"), map.get("value")) { + if let (Some(k_str), Some(v_str)) = (k.as_str(), v.as_str()) { + res.insert(k_str.to_string(), v_str.to_string()); + } else { + return Err(ArrowError::ParseError(format!( + "Field '{attribute}' must have map value of string type" + ))); + } + } else { + return Err(ArrowError::ParseError(format!( + "Field '{attribute}' lacks map keys named \"key\" or \"value\"" + ))); + } + } + _ => { + return Err(ArrowError::ParseError(format!( + "Field '{attribute}' contains non-object key-value pair" + ))); + } + } + } + Ok(res) + } + Value::Object(values) => { + let mut res: HashMap = HashMap::default(); + for (k, v) in values { + if let Some(str_value) = v.as_str() { + res.insert(k.clone(), str_value.to_string()); + } else { + return Err(ArrowError::ParseError(format!( + "Field '{attribute}' contains non-string value for key {k}" + ))); + } + } + Ok(res) + } + _ => Err(ArrowError::ParseError(format!( + "Field `{attribute}` is not json array" + ))), + } +} + /// Generate a JSON representation of the `Field`. pub fn field_to_json(field: &Field) -> serde_json::Value { let children: Vec = match field.data_type() { @@ -616,4 +646,132 @@ mod tests { assert_eq!(expected, dt); } + + /// JSON of a dictionary-encoded `fixedsizebinary(16)` field, with the given + /// `dictionary.metadata` and field `metadata` attributes (if any) + fn uuid_dictionary_field_json( + dictionary_metadata: Option<&str>, + metadata: Option<&str>, + ) -> Value { + let mut json = serde_json::json!({ + "name": "uuids", + "nullable": true, + "type": {"name": "fixedsizebinary", "byteWidth": 16}, + "children": [], + "dictionary": { + "id": 0, + "indexType": {"name": "int", "isSigned": true, "bitWidth": 32}, + "isOrdered": false + } + }); + if let Some(dictionary_metadata) = dictionary_metadata { + json["dictionary"]["metadata"] = serde_json::from_str(dictionary_metadata).unwrap(); + } + if let Some(metadata) = metadata { + json["metadata"] = serde_json::from_str(metadata).unwrap(); + } + json + } + + fn uuid_dictionary_type() -> DataType { + DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::FixedSizeBinary(16)), + ) + } + + /// Read a dictionary field with the given `dictionary.metadata` and field + /// `metadata` (as `[{"key": ..., "value": ...}]` lists) + fn read_uuid_dictionary_field( + dictionary_metadata: Option<&[(&str, &str)]>, + metadata: Option<&[(&str, &str)]>, + ) -> Field { + let to_json = |pairs: &[(&str, &str)]| { + let pairs = pairs + .iter() + .map(|(k, v)| serde_json::json!({"key": k, "value": v})) + .collect::>(); + serde_json::to_string(&pairs).unwrap() + }; + let json = uuid_dictionary_field_json( + dictionary_metadata.map(to_json).as_deref(), + metadata.map(to_json).as_deref(), + ); + let field = field_from_json(&json).unwrap(); + assert_eq!(field.data_type(), &uuid_dictionary_type()); + field + } + + #[test] + fn parse_dictionary_value_metadata_from_json() { + let name = EXTENSION_TYPE_NAME_KEY; + let ext_metadata = "ARROW:extension:metadata"; + let expected = |pairs: &[(&str, &str)]| { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>() + }; + + // copied to the field metadata + let field = read_uuid_dictionary_field(Some(&[(name, "arrow.uuid")]), None); + assert_eq!(field.metadata(), &expected(&[(name, "arrow.uuid")])); + + // without replacing the keys of the field metadata + let field = read_uuid_dictionary_field( + Some(&[ + (name, "arrow.uuid"), + (ext_metadata, "value"), + ("k", "value"), + ]), + Some(&[(ext_metadata, "field"), ("k", "field")]), + ); + assert_eq!( + field.metadata(), + &expected(&[ + (name, "arrow.uuid"), + (ext_metadata, "field"), + ("k", "field") + ]) + ); + + // dropped if the field has its own extension type + let field = read_uuid_dictionary_field( + Some(&[(name, "arrow.uuid"), ("k", "value")]), + Some(&[(name, "example.uuid_dictionary")]), + ); + assert_eq!( + field.metadata(), + &expected(&[(name, "example.uuid_dictionary")]) + ); + + // empty + let field = read_uuid_dictionary_field(Some(&[]), Some(&[("k", "field")])); + assert_eq!( + field, + read_uuid_dictionary_field(None, Some(&[("k", "field")])) + ); + } + + #[test] + fn parse_invalid_dictionary_value_metadata_from_json() { + for (metadata, error) in [ + ( + r#""arrow.uuid""#, + "Field `dictionary.metadata` is not json array", + ), + ( + r#"[{"key": "ARROW:extension:name"}]"#, + "Field 'dictionary.metadata' must have exact two entries for each key-value map", + ), + ( + r#"[{"key": "ARROW:extension:name", "value": 1}]"#, + "Field 'dictionary.metadata' must have map value of string type", + ), + ] { + let value = uuid_dictionary_field_json(Some(metadata), None); + let err = field_from_json(&value).unwrap_err(); + assert_eq!(err.to_string(), format!("Parser error: {error}")); + } + } } diff --git a/arrow-integration-test/src/lib.rs b/arrow-integration-test/src/lib.rs index 12f4f2b5d5bb..597356aed5ee 100644 --- a/arrow-integration-test/src/lib.rs +++ b/arrow-integration-test/src/lib.rs @@ -149,6 +149,10 @@ impl From<&Field> for ArrowJsonField { } /// Represents a dictionary-encoded field in the Arrow JSON format +/// +/// The optional `metadata` of the dictionary value type is not kept, so it is +/// not seen when converting an [`ArrowJsonField`] to a [`Field`]. Use +/// [`field_from_json`] to read it. #[derive(Deserialize, Serialize, Debug)] pub struct ArrowJsonFieldDictionary { /// A unique identifier for the dictionary diff --git a/arrow-ipc/src/convert.rs b/arrow-ipc/src/convert.rs index 9affcd6b4af6..8c4770961633 100644 --- a/arrow-ipc/src/convert.rs +++ b/arrow-ipc/src/convert.rs @@ -159,8 +159,13 @@ pub fn schema_to_fb_offset<'a>( /// Convert an IPC Field to Arrow Field /// +/// See [value type metadata] for how the metadata of a dictionary's value +/// type is mapped to [`Field`] metadata. +/// /// This panics on malformed input; every reader path uses the fallible /// conversion instead. kept for backwards compatibility only. +/// +/// [value type metadata]: crate::reader#dictionary-value-type-metadata impl From> for Field { fn from(field: crate::Field) -> Field { try_field_from(field).expect("invalid IPC field") @@ -168,6 +173,11 @@ impl From> for Field { } /// Convert an IPC Field to Arrow Field +/// +/// See [value type metadata] for how `DictionaryEncoding.custom_metadata` +/// is mapped. +/// +/// [value type metadata]: crate::reader#dictionary-value-type-metadata fn try_field_from(field: crate::Field) -> Result { let arrow_field = if let Some(dictionary) = field.dictionary() { #[expect(deprecated)] @@ -195,6 +205,21 @@ fn try_field_from(field: crate::Field) -> Result { } } + // `DataType::Dictionary` cannot carry metadata for its value type + // (https://github.com/apache/arrow-rs/issues/7982), so copy it to the field + // unless the field has its own extension type + if let Some(value_metadata) = field.dictionary().and_then(|d| d.custom_metadata()) + && !metadata_map.contains_key(extension::EXTENSION_TYPE_NAME_KEY) + { + for kv in value_metadata { + if let (Some(k), Some(v)) = (kv.key(), kv.value()) { + metadata_map + .entry(k.to_string()) + .or_insert_with(|| v.to_string()); + } + } + } + Ok(arrow_field.with_metadata(metadata_map)) } @@ -208,6 +233,11 @@ pub fn fb_to_schema(fb: crate::Schema) -> Schema { /// /// Returns an error on schema messages that the flatbuffer /// verifier accepts but that are not valid Arrow. +/// +/// See [value type metadata] for how the metadata of a dictionary's value +/// type is mapped to [`Field`] metadata. +/// +/// [value type metadata]: crate::reader#dictionary-value-type-metadata pub fn try_fb_to_schema(fb: crate::Schema) -> Result { let mut fields: Vec = vec![]; let c_fields = fb @@ -245,6 +275,11 @@ pub fn try_fb_to_schema(fb: crate::Schema) -> Result { } /// Try deserialize flat buffer format bytes into a schema +/// +/// See [value type metadata] for how the metadata of a dictionary's value +/// type is mapped to [`Field`] metadata. +/// +/// [value type metadata]: crate::reader#dictionary-value-type-metadata pub fn try_schema_from_flatbuffer_bytes(bytes: &[u8]) -> Result { if let Ok(ipc) = crate::root_as_message(bytes) { if let Some(schema) = ipc.header_as_schema() { @@ -262,6 +297,11 @@ pub fn try_schema_from_flatbuffer_bytes(bytes: &[u8]) -> Result Result { // There are two protocol types: https://issues.apache.org/jira/browse/ARROW-6313 // The original protocol is: @@ -1018,6 +1058,13 @@ pub(crate) fn get_fb_field_type<'a>( } /// Create an IPC dictionary encoding +/// +/// `DictionaryEncoding.custom_metadata` (metadata for the dictionary value +/// type) is never written: [`DataType::Dictionary`] cannot carry metadata for +/// its value type (see ), so +/// arrow-rs cannot distinguish metadata of the dictionary value type from +/// metadata of the field's own type, and writes the metadata of a dictionary +/// field to `Field.custom_metadata`. pub(crate) fn get_fb_dictionary<'a>( index_type: &DataType, dict_id: i64, @@ -1091,6 +1138,136 @@ impl MessageBuffer { } } +/// Helpers to build IPC schemas that the arrow-rs writer cannot produce +#[cfg(test)] +pub(crate) mod test_util { + use super::*; + + pub(crate) const UUID: &str = "arrow.uuid"; + + /// `Dictionary(Int32, FixedSizeBinary(16))`, the type of the fields built + /// by [`fb_dictionary_field`] + pub(crate) fn uuid_dictionary_type() -> DataType { + Dictionary(Box::new(Int32), Box::new(FixedSizeBinary(16))) + } + + /// Push `pairs` as a vector of IPC `KeyValue`s. + fn fb_key_values<'a>( + fbb: &mut FlatBufferBuilder<'a>, + pairs: &[(&str, &str)], + ) -> WIPOffset>>> { + let pairs = pairs + .iter() + .map(|(k, v)| { + let key = fbb.create_string(k); + let value = fbb.create_string(v); + let mut kv = crate::KeyValueBuilder::new(fbb); + kv.add_key(key); + kv.add_value(value); + kv.finish() + }) + .collect::>(); + fbb.create_vector(&pairs) + } + + /// Push an IPC field `name` of type [`uuid_dictionary_type`] with + /// dictionary id `id`, the given `DictionaryEncoding.custom_metadata` + /// (omitted if `None`) and `Field.custom_metadata`. + pub(crate) fn fb_dictionary_field<'a>( + fbb: &mut FlatBufferBuilder<'a>, + name: &str, + id: i64, + value_metadata: Option<&[(&str, &str)]>, + field_metadata: &[(&str, &str)], + ) -> WIPOffset> { + let name = fbb.create_string(name); + let value_type = { + let mut b = crate::FixedSizeBinaryBuilder::new(fbb); + b.add_byteWidth(16); + b.finish().as_union_value() + }; + let index_type = { + let mut b = crate::IntBuilder::new(fbb); + b.add_bitWidth(32); + b.add_is_signed(true); + b.finish() + }; + let value_metadata = value_metadata.map(|m| fb_key_values(fbb, m)); + let dictionary = crate::DictionaryEncoding::create( + fbb, + &crate::DictionaryEncodingArgs { + id, + indexType: Some(index_type), + isOrdered: false, + dictionaryKind: crate::DictionaryKind::DenseArray, + custom_metadata: value_metadata, + }, + ); + let field_metadata = + (!field_metadata.is_empty()).then(|| fb_key_values(fbb, field_metadata)); + let mut b = crate::FieldBuilder::new(fbb); + b.add_name(name); + b.add_nullable(true); + b.add_type_type(crate::Type::FixedSizeBinary); + b.add_type_(value_type); + b.add_dictionary(dictionary); + if let Some(field_metadata) = field_metadata { + b.add_custom_metadata(field_metadata); + } + b.finish() + } + + /// Push an IPC field `name` of type `type_type` (a list or a struct) with + /// the given children. + pub(crate) fn fb_nested_field<'a>( + fbb: &mut FlatBufferBuilder<'a>, + name: &str, + type_type: crate::Type, + children: &[WIPOffset>], + ) -> WIPOffset> { + let name = fbb.create_string(name); + let type_ = match type_type { + crate::Type::List => crate::ListBuilder::new(fbb).finish().as_union_value(), + crate::Type::Struct_ => crate::Struct_Builder::new(fbb).finish().as_union_value(), + _ => unimplemented!("{type_type:?}"), + }; + let children = fbb.create_vector(children); + let mut b = crate::FieldBuilder::new(fbb); + b.add_name(name); + b.add_nullable(true); + b.add_type_type(type_type); + b.add_type_(type_); + b.add_children(children); + b.finish() + } + + /// Push an IPC schema of `fields`. + pub(crate) fn fb_schema<'a>( + fbb: &mut FlatBufferBuilder<'a>, + fields: &[WIPOffset>], + ) -> WIPOffset> { + let fields = fbb.create_vector(fields); + let mut b = crate::SchemaBuilder::new(fbb); + b.add_fields(fields); + b.finish() + } + + /// Finish `fbb` as an IPC `Message` whose header is a schema of `fields`. + pub(crate) fn fb_schema_message<'a>( + mut fbb: FlatBufferBuilder<'a>, + fields: &[WIPOffset>], + ) -> Vec { + let schema = fb_schema(&mut fbb, fields); + let mut message = crate::MessageBuilder::new(&mut fbb); + message.add_version(crate::MetadataVersion::V5); + message.add_header_type(crate::MessageHeader::Schema); + message.add_header(schema.as_union_value()); + let message = message.finish(); + fbb.finish(message, None); + fbb.finished_data().to_vec() + } +} + #[cfg(test)] mod tests { use super::*; @@ -1432,4 +1609,131 @@ mod tests { assert!(ipc.custom_metadata().is_none()); assert!(ipc2.custom_metadata().is_none()); } + + use super::test_util::*; + use extension::{EXTENSION_TYPE_METADATA_KEY, EXTENSION_TYPE_NAME_KEY}; + + /// Read the single field of a schema message holding the given dictionary + /// field + fn read_dictionary_field( + value_metadata: Option<&[(&str, &str)]>, + field_metadata: &[(&str, &str)], + ) -> Field { + let mut fbb = FlatBufferBuilder::new(); + let field = fb_dictionary_field(&mut fbb, "uuids", 0, value_metadata, field_metadata); + let bytes = fb_schema_message(fbb, &[field]); + let schema = try_schema_from_flatbuffer_bytes(&bytes).unwrap(); + let field = schema.field(0).clone(); + assert_eq!(field.data_type(), &uuid_dictionary_type()); + field + } + + #[test] + fn dictionary_value_metadata_is_copied_to_field() { + let field = read_dictionary_field( + Some(&[ + (EXTENSION_TYPE_NAME_KEY, UUID), + (EXTENSION_TYPE_METADATA_KEY, ""), + ]), + &[], + ); + assert_eq!( + field.metadata(), + &Metadata::from([ + (EXTENSION_TYPE_NAME_KEY, UUID), + (EXTENSION_TYPE_METADATA_KEY, ""), + ]) + ); + } + + #[test] + fn dictionary_value_metadata_does_not_override_field_metadata() { + let field = read_dictionary_field( + Some(&[ + (EXTENSION_TYPE_NAME_KEY, UUID), + (EXTENSION_TYPE_METADATA_KEY, "value"), + ("k", "value"), + ]), + &[(EXTENSION_TYPE_METADATA_KEY, "field"), ("k", "field")], + ); + assert_eq!( + field.metadata(), + &Metadata::from([ + (EXTENSION_TYPE_NAME_KEY, UUID), + (EXTENSION_TYPE_METADATA_KEY, "field"), + ("k", "field"), + ]) + ); + } + + #[test] + fn dictionary_value_metadata_is_dropped_for_field_extension_type() { + // an extension type with dictionary storage, whose dictionary values + // are themselves an extension type + let field_metadata = [(EXTENSION_TYPE_NAME_KEY, "example.uuid_dictionary")]; + let field = read_dictionary_field( + Some(&[(EXTENSION_TYPE_NAME_KEY, UUID), ("k", "value")]), + &field_metadata, + ); + assert_eq!(field.metadata(), &Metadata::from(field_metadata)); + } + + #[test] + fn empty_dictionary_value_metadata() { + let field_metadata = [("k", "field")]; + let field = read_dictionary_field(Some(&[]), &field_metadata); + assert_eq!(field, read_dictionary_field(None, &field_metadata)); + assert_eq!(field.metadata(), &Metadata::from(field_metadata)); + } + + #[test] + fn nested_dictionary_value_metadata_is_copied_to_field() { + let mut fbb = FlatBufferBuilder::new(); + let value_metadata: &[_] = &[(EXTENSION_TYPE_NAME_KEY, UUID)]; + let item = fb_dictionary_field(&mut fbb, "item", 0, Some(value_metadata), &[]); + let list = fb_nested_field(&mut fbb, "list", crate::Type::List, &[item]); + let child = fb_dictionary_field(&mut fbb, "child", 1, Some(value_metadata), &[]); + let r#struct = fb_nested_field(&mut fbb, "struct", crate::Type::Struct_, &[child]); + let bytes = fb_schema_message(fbb, &[list, r#struct]); + + let schema = try_schema_from_flatbuffer_bytes(&bytes).unwrap(); + let uuid_field = |name| { + #[expect(deprecated)] + Field::new_dict(name, uuid_dictionary_type(), true, 0, false) + .with_metadata([(EXTENSION_TYPE_NAME_KEY, UUID)]) + }; + let expected = Schema::new(vec![ + Field::new_list("list", uuid_field("item"), true), + Field::new_struct("struct", vec![uuid_field("child")], true), + ]); + assert_eq!(schema, expected); + } + + #[test] + fn non_utf8_dictionary_value_metadata_is_rejected() { + // the flatbuffer verifier checks that the strings of + // `DictionaryEncoding.custom_metadata` are valid UTF-8 + let marker = "not_utf8_marker"; + let mut fbb = FlatBufferBuilder::new(); + let field = fb_dictionary_field( + &mut fbb, + "uuids", + 0, + Some(&[(EXTENSION_TYPE_NAME_KEY, marker)]), + &[], + ); + let mut bytes = fb_schema_message(fbb, &[field]); + assert!(try_schema_from_flatbuffer_bytes(&bytes).is_ok()); + + let pos = bytes + .windows(marker.len()) + .position(|w| w == marker.as_bytes()) + .unwrap(); + bytes[pos] = 0xff; + let err = try_schema_from_flatbuffer_bytes(&bytes).unwrap_err(); + assert_eq!( + err.to_string(), + "Parser error: Unable to get root as message" + ); + } } diff --git a/arrow-ipc/src/gen/Schema.rs b/arrow-ipc/src/gen/Schema.rs index f15072b517cb..f1606ce37058 100644 --- a/arrow-ipc/src/gen/Schema.rs +++ b/arrow-ipc/src/gen/Schema.rs @@ -4129,6 +4129,7 @@ impl<'a> DictionaryEncoding<'a> { pub const VT_INDEXTYPE: flatbuffers::VOffsetT = 6; pub const VT_ISORDERED: flatbuffers::VOffsetT = 8; pub const VT_DICTIONARYKIND: flatbuffers::VOffsetT = 10; + pub const VT_CUSTOM_METADATA: flatbuffers::VOffsetT = 12; #[inline] pub unsafe fn init_from_table(table: flatbuffers::Table<'a>) -> Self { @@ -4141,6 +4142,9 @@ impl<'a> DictionaryEncoding<'a> { ) -> flatbuffers::WIPOffset> { let mut builder = DictionaryEncodingBuilder::new(_fbb); builder.add_id(args.id); + if let Some(x) = args.custom_metadata { + builder.add_custom_metadata(x); + } if let Some(x) = args.indexType { builder.add_indexType(x); } @@ -4207,6 +4211,20 @@ impl<'a> DictionaryEncoding<'a> { .unwrap() } } + /// Metadata for the dictionary value type, e.g. the extension type of the values + #[inline] + pub fn custom_metadata( + &self, + ) -> Option>>> { + // Safety: + // Created from valid Table for this object + // which contains a valid value in this slot + unsafe { + self._tab.get::>, + >>(DictionaryEncoding::VT_CUSTOM_METADATA, None) + } + } } impl flatbuffers::Verifiable for DictionaryEncoding<'_> { @@ -4225,6 +4243,9 @@ impl flatbuffers::Verifiable for DictionaryEncoding<'_> { )? .visit_field::("isOrdered", Self::VT_ISORDERED, false)? .visit_field::("dictionaryKind", Self::VT_DICTIONARYKIND, false)? + .visit_field::>, + >>("custom_metadata", Self::VT_CUSTOM_METADATA, false)? .finish(); Ok(()) } @@ -4234,6 +4255,9 @@ pub struct DictionaryEncodingArgs<'a> { pub indexType: Option>>, pub isOrdered: bool, pub dictionaryKind: DictionaryKind, + pub custom_metadata: Option< + flatbuffers::WIPOffset>>>, + >, } impl<'a> Default for DictionaryEncodingArgs<'a> { #[inline] @@ -4243,6 +4267,7 @@ impl<'a> Default for DictionaryEncodingArgs<'a> { indexType: None, isOrdered: false, dictionaryKind: DictionaryKind::DenseArray, + custom_metadata: None, } } } @@ -4277,6 +4302,18 @@ impl<'a: 'b, 'b, A: flatbuffers::Allocator + 'a> DictionaryEncodingBuilder<'a, ' ); } #[inline] + pub fn add_custom_metadata( + &mut self, + custom_metadata: flatbuffers::WIPOffset< + flatbuffers::Vector<'b, flatbuffers::ForwardsUOffset>>, + >, + ) { + self.fbb_.push_slot_always::>( + DictionaryEncoding::VT_CUSTOM_METADATA, + custom_metadata, + ); + } + #[inline] pub fn new( _fbb: &'b mut flatbuffers::FlatBufferBuilder<'a, A>, ) -> DictionaryEncodingBuilder<'a, 'b, A> { @@ -4300,6 +4337,7 @@ impl core::fmt::Debug for DictionaryEncoding<'_> { ds.field("indexType", &self.indexType()); ds.field("isOrdered", &self.isOrdered()); ds.field("dictionaryKind", &self.dictionaryKind()); + ds.field("custom_metadata", &self.custom_metadata()); ds.finish() } } diff --git a/arrow-ipc/src/reader.rs b/arrow-ipc/src/reader.rs index f891b6dd3c85..f56aeecb9a5f 100644 --- a/arrow-ipc/src/reader.rs +++ b/arrow-ipc/src/reader.rs @@ -23,6 +23,37 @@ //! however the [`FileReader`] expects a reader that supports [`Seek`]ing //! //! [`Seek`]: std::io::Seek +//! +//! # Dictionary value type metadata +//! +//! Since version 1.6 of the IPC format, the `DictionaryEncoding` of a +//! dictionary-encoded field can hold `custom_metadata` for the dictionary +//! *value* type, for example the `ARROW:extension:name` of a dictionary whose +//! values are an extension type, while `Field.custom_metadata` describes the +//! field's own type, for example an extension type whose storage type is a +//! dictionary. +//! +//! [`DataType::Dictionary`] cannot carry metadata for its value type (see +//! [#7982](https://github.com/apache/arrow-rs/issues/7982)), so until it can, +//! arrow-rs maps the value type metadata to the metadata of the [`Field`]: +//! +//! * If the field's own metadata has no `ARROW:extension:name`, the value type +//! metadata is copied into the field's metadata. Keys already present in the +//! field's metadata take precedence. +//! * Otherwise the value type metadata is dropped. +//! +//! This applies to the schemas read by [`FileReader`], [`StreamReader`] and +//! the schema conversions in [`crate::convert`]. The C Data Interface import +//! (`Field::try_from(&FFI_ArrowSchema)`) maps the metadata of the `dictionary` +//! of an imported schema the same way. It keeps metadata such as an +//! `arrow.uuid` extension type on dictionary-encoded fields, where writers +//! that predate version 1.6 put it. +//! +//! arrow-rs never writes `DictionaryEncoding.custom_metadata`: the metadata of +//! a dictionary-encoded field is always written to `Field.custom_metadata`, +//! the position used before version 1.6, which other implementations read as +//! metadata of the field's own type (for example an extension type whose +//! storage type is a dictionary). mod stream; pub use stream::*; @@ -4296,4 +4327,134 @@ mod tests { "unexpected error: {err}" ); } + + /// A batch of a uuid dictionary column, without extension metadata + fn uuid_dictionary_batch() -> RecordBatch { + use crate::convert::test_util::uuid_dictionary_type; + let values = + FixedSizeBinaryArray::try_from_iter([[1u8; 16], [2u8; 16]].into_iter()).unwrap(); + let keys = Int32Array::from(vec![Some(1), None, Some(0), Some(1)]); + let array = DictionaryArray::::try_new(keys, Arc::new(values)).unwrap(); + let schema = Schema::new(vec![Field::new("uuids", uuid_dictionary_type(), true)]); + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap() + } + + /// The schema of [`uuid_dictionary_batch`], with `arrow.uuid` as the + /// dictionary value type + fn fb_uuid_dictionary_schema<'a>( + fbb: &mut flatbuffers::FlatBufferBuilder<'a>, + ) -> Vec>> { + use crate::convert::test_util::{UUID, fb_dictionary_field}; + let value_metadata: &[_] = &[(extension::EXTENSION_TYPE_NAME_KEY, UUID)]; + vec![fb_dictionary_field( + fbb, + "uuids", + 0, + Some(value_metadata), + &[], + )] + } + + /// Assert that `schema` and `batches` are [`uuid_dictionary_batch`], with the + /// dictionary value type metadata copied to the field + fn assert_uuid_dictionary_batch(schema: SchemaRef, batches: Vec) { + let batch = uuid_dictionary_batch(); + let field = batch.schema().field(0).clone().with_metadata([( + extension::EXTENSION_TYPE_NAME_KEY, + crate::convert::test_util::UUID, + )]); + let expected = Arc::new(Schema::new(vec![field])); + assert_eq!(schema, expected); + assert_eq!(batches, vec![batch.with_schema(expected).unwrap()]); + } + + #[test] + fn test_stream_reader_dictionary_value_metadata() { + use crate::convert::test_util::fb_schema_message; + + let batch = uuid_dictionary_batch(); + let mut stream = vec![]; + { + let mut writer = + crate::writer::StreamWriter::try_new(&mut stream, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // replace the schema message with one whose dictionary encoding has + // metadata for the dictionary value type + assert_eq!(stream[..4], CONTINUATION_MARKER); + let schema_len = i32::from_le_bytes(stream[4..8].try_into().unwrap()) as usize; + let mut fbb = flatbuffers::FlatBufferBuilder::new(); + let fields = fb_uuid_dictionary_schema(&mut fbb); + let mut message = fb_schema_message(fbb, &fields); + message.resize(message.len().next_multiple_of(8), 0); + let mut new_stream = CONTINUATION_MARKER.to_vec(); + new_stream.extend_from_slice(&(message.len() as i32).to_le_bytes()); + new_stream.extend_from_slice(&message); + new_stream.extend_from_slice(&stream[8 + schema_len..]); + + let reader = StreamReader::try_new(std::io::Cursor::new(new_stream), None).unwrap(); + let schema = reader.schema(); + let batches = reader.collect::, _>>().unwrap(); + assert_uuid_dictionary_batch(schema, batches); + } + + #[test] + fn test_file_reader_dictionary_value_metadata() { + use crate::convert::test_util::fb_schema; + + let batch = uuid_dictionary_batch(); + let mut file = vec![]; + { + let mut writer = + crate::writer::FileWriter::try_new(&mut file, &batch.schema()).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + + // replace the schema in the footer with one whose dictionary encoding + // has metadata for the dictionary value type + let (data, trailer) = file.split_at(file.len() - 10); + assert_eq!(&trailer[4..], crate::ARROW_MAGIC); + let footer_len = i32::from_le_bytes(trailer[..4].try_into().unwrap()) as usize; + let (data, footer) = data.split_at(data.len() - footer_len); + let footer = root_as_footer(footer).unwrap(); + let dictionaries = footer + .dictionaries() + .unwrap() + .iter() + .copied() + .collect::>(); + let record_batches = footer + .recordBatches() + .unwrap() + .iter() + .copied() + .collect::>(); + + let mut fbb = flatbuffers::FlatBufferBuilder::new(); + let fields = fb_uuid_dictionary_schema(&mut fbb); + let schema = fb_schema(&mut fbb, &fields); + let dictionaries = fbb.create_vector(&dictionaries); + let record_batches = fbb.create_vector(&record_batches); + let mut new_footer = crate::FooterBuilder::new(&mut fbb); + new_footer.add_version(footer.version()); + new_footer.add_schema(schema); + new_footer.add_dictionaries(dictionaries); + new_footer.add_recordBatches(record_batches); + let new_footer = new_footer.finish(); + fbb.finish(new_footer, None); + let new_footer = fbb.finished_data(); + + let mut new_file = data.to_vec(); + new_file.extend_from_slice(new_footer); + new_file.extend_from_slice(&(new_footer.len() as i32).to_le_bytes()); + new_file.extend_from_slice(&trailer[4..]); + + let reader = FileReader::try_new(std::io::Cursor::new(new_file), None).unwrap(); + let schema = reader.schema(); + let batches = reader.collect::, _>>().unwrap(); + assert_uuid_dictionary_batch(schema, batches); + } } diff --git a/arrow-schema/src/ffi.rs b/arrow-schema/src/ffi.rs index 85fb32728593..574b812031c6 100644 --- a/arrow-schema/src/ffi.rs +++ b/arrow-schema/src/ffi.rs @@ -34,6 +34,7 @@ //! assert_eq!(schema, back); //! ``` +use crate::extension::EXTENSION_TYPE_NAME_KEY; use crate::{ ArrowError, DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, UnionFields, UnionMode, }; @@ -716,14 +717,41 @@ impl TryFrom<&FFI_ArrowSchema> for DataType { } } +/// Import a [`Field`] from an [`FFI_ArrowSchema`] +/// +/// # Dictionary value type metadata +/// +/// The metadata of the `dictionary` of a dictionary-encoded schema describes +/// the dictionary *value* type, for example the extension type of a dictionary +/// whose values are an extension type. [`DataType::Dictionary`] cannot carry +/// metadata for its value type (see +/// ), so until it can it is +/// mapped to the metadata of the [`Field`], as the arrow-ipc readers map the +/// metadata of the dictionary value type in IPC schemas: +/// +/// * If the schema's own metadata has no `ARROW:extension:name`, the metadata +/// of its `dictionary` is copied into the field's metadata. Keys already +/// present in the schema's own metadata take precedence. +/// * Otherwise the metadata of its `dictionary` is dropped. +/// +/// Exporting a [`Field`] writes all its metadata to the exported schema itself, +/// not to its `dictionary`. impl TryFrom<&FFI_ArrowSchema> for Field { type Error = ArrowError; fn try_from(c_schema: &FFI_ArrowSchema) -> Result { let dtype = DataType::try_from(c_schema)?; + let mut metadata = c_schema.metadata()?; + if let Some(dict_schema) = c_schema.dictionary() + && !metadata.contains_key(EXTENSION_TYPE_NAME_KEY) + { + for (k, v) in dict_schema.metadata()? { + metadata.entry(k).or_insert(v); + } + } let field = Field::new(c_schema.name().unwrap_or(""), dtype, c_schema.nullable()) .with_dict_is_ordered(c_schema.dictionary_ordered()) - .with_metadata(c_schema.metadata()?); + .with_metadata(metadata); Ok(field) } } @@ -1113,6 +1141,118 @@ mod tests { } } + /// Build a nullable schema `name` of a dictionary of `fixed_size_binary(16)` + /// with int32 indices, with the given metadata on the schema and on its + /// dictionary. + fn uuid_dictionary_schema( + name: &str, + value_metadata: &[(&str, &str)], + metadata: &[(&str, &str)], + ) -> FFI_ArrowSchema { + // SAFETY: the schemas were constructed by this crate via try_new + unsafe { + let dictionary = FFI_ArrowSchema::try_new("w:16", vec![], None) + .unwrap() + .with_metadata(value_metadata.iter().copied()) + .unwrap(); + FFI_ArrowSchema::try_new("i", vec![], Some(dictionary)) + .unwrap() + .with_name(name) + .unwrap() + .with_flags(Flags::NULLABLE) + .unwrap() + .with_metadata(metadata.iter().copied()) + .unwrap() + } + } + + /// The field imported from [`uuid_dictionary_schema`] with `metadata` + fn uuid_dictionary_field(name: &str, metadata: &[(&str, &str)]) -> Field { + let data_type = DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::FixedSizeBinary(16)), + ); + Field::new(name, data_type, true).with_metadata( + metadata + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>(), + ) + } + + #[test] + fn test_import_dictionary_value_metadata() { + const NAME: &str = EXTENSION_TYPE_NAME_KEY; + const METADATA: &str = crate::extension::EXTENSION_TYPE_METADATA_KEY; + let import = |value_metadata, metadata| { + Field::try_from(&uuid_dictionary_schema("uuids", value_metadata, metadata)).unwrap() + }; + + // copied to the field + assert_eq!( + import(&[(NAME, "arrow.uuid"), (METADATA, "")], &[]), + uuid_dictionary_field("uuids", &[(NAME, "arrow.uuid"), (METADATA, "")]) + ); + + // without replacing the keys of the field's own metadata + assert_eq!( + import( + &[(NAME, "arrow.uuid"), (METADATA, "value"), ("k", "value")], + &[(METADATA, "field"), ("k", "field")] + ), + uuid_dictionary_field( + "uuids", + &[(NAME, "arrow.uuid"), (METADATA, "field"), ("k", "field")] + ) + ); + + // dropped if the field has its own extension type + assert_eq!( + import( + &[(NAME, "arrow.uuid"), ("k", "value")], + &[(NAME, "example.uuid_dictionary")] + ), + uuid_dictionary_field("uuids", &[(NAME, "example.uuid_dictionary")]) + ); + } + + #[test] + fn test_import_nested_dictionary_value_metadata() { + let value_metadata = [(EXTENSION_TYPE_NAME_KEY, "arrow.uuid")]; + let list = FFI_ArrowSchema::try_new( + "+l", + vec![uuid_dictionary_schema("item", &value_metadata, &[])], + None, + ) + .unwrap() + .with_name("list") + .unwrap(); + let r#struct = FFI_ArrowSchema::try_new( + "+s", + vec![uuid_dictionary_schema("child", &value_metadata, &[])], + None, + ) + .unwrap() + .with_name("struct") + .unwrap(); + let c_schema = FFI_ArrowSchema::try_new("+s", vec![list, r#struct], None).unwrap(); + + let schema = Schema::try_from(&c_schema).unwrap(); + let expected = Schema::new(vec![ + Field::new_list( + "list", + uuid_dictionary_field("item", &value_metadata), + false, + ), + Field::new_struct( + "struct", + vec![uuid_dictionary_field("child", &value_metadata)], + false, + ), + ]); + assert_eq!(schema, expected); + } + #[test] fn test_name_with_null_byte() { let schema = FFI_ArrowSchema::try_new("i", vec![], None).unwrap(); diff --git a/format/Schema.fbs b/format/Schema.fbs index 19e1633e809c..38a92f906f3f 100644 --- a/format/Schema.fbs +++ b/format/Schema.fbs @@ -25,6 +25,7 @@ /// Version 1.4 - Add BinaryView, Utf8View, variadicBufferCounts, ListView, and /// LargeListView. /// Version 1.5 - Add 32-bit and 64-bit as allowed bit widths for Decimal +/// Version 1.6 - Add custom_metadata to DictionaryEncoding. namespace org.apache.arrow.flatbuf; @@ -511,6 +512,9 @@ table DictionaryEncoding { isOrdered: bool; dictionaryKind: DictionaryKind; + + /// Metadata for the dictionary value type, e.g. the extension type of the values + custom_metadata: [ KeyValue ]; } /// ----------------------------------------------------------------------