Add IpcSchemaEncoder, deprecate ipc schema functions, Fix IPC not respecting not preserving dict ID

arrow-flight/src/lib.rs

@@ -137,7 +137,9 @@ pub struct IpcMessage(pub Bytes);
 
 fn flight_schema_as_encoded_data(arrow_schema: &Schema, options: &IpcWriteOptions) -> EncodedData {
     let data_gen = writer::IpcDataGenerator::default();
-    data_gen.schema_to_bytes(arrow_schema, options)
+    let mut dict_tracker =
+        writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());
+    data_gen.schema_to_bytes_with_dictionary_tracker(arrow_schema, &mut dict_tracker, options)
 }
 
 fn flight_schema_as_flatbuffer(schema: &Schema, options: &IpcWriteOptions) -> IpcMessage {

arrow-ipc/src/convert.rs

@@ -27,13 +27,63 @@ use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 
+use crate::writer::DictionaryTracker;
 use crate::{size_prefixed_root_as_message, KeyValue, Message, CONTINUATION_MARKER};
 use DataType::*;
 
+pub struct IpcSchemaConverter<'a> {
+    dictionary_tracker: &'a mut DictionaryTracker,
+}
+
+impl IpcSchemaConverter<'_> {
+    pub fn new(dictionary_tracker: &mut DictionaryTracker) -> IpcSchemaConverter<'_> {
+        IpcSchemaConverter { dictionary_tracker }
+    }
+
+    /// Serialize a schema in IPC format
+    pub fn schema_to_fb<'a>(&mut self, schema: &'a Schema) -> FlatBufferBuilder<'a> {
+        let mut fbb = FlatBufferBuilder::new();
+
+        let root = self.schema_to_fb_offset(&mut fbb, schema);
+
+        fbb.finish(root, None);
+
+        fbb
+    }
+
+    pub fn schema_to_fb_offset<'a>(
+        &mut self,
+        fbb: &mut FlatBufferBuilder<'a>,
+        schema: &Schema,
+    ) -> WIPOffset<crate::Schema<'a>> {
+        let fields = schema
+            .fields()
+            .iter()
+            .map(|field| build_field(fbb, &mut Some(self.dictionary_tracker), field))
+            .collect::<Vec<_>>();
+        let fb_field_list = fbb.create_vector(&fields);
+
+        let fb_metadata_list =
+            (!schema.metadata().is_empty()).then(|| metadata_to_fb(fbb, schema.metadata()));
+
+        let mut builder = crate::SchemaBuilder::new(fbb);
+        builder.add_fields(fb_field_list);
+        if let Some(fb_metadata_list) = fb_metadata_list {
+            builder.add_custom_metadata(fb_metadata_list);
+        }
+        builder.finish()
+    }
+}
+
 /// Serialize a schema in IPC format
-pub fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder {
+#[deprecated(
+    since = "54.0.0",
+    note = "Use `IpcSchemaConverter` instead. This function will be removed in the next release."
+)]
+pub fn schema_to_fb(schema: &Schema) -> FlatBufferBuilder<'_> {
     let mut fbb = FlatBufferBuilder::new();
 
+    #[allow(deprecated)]
     let root = schema_to_fb_offset(&mut fbb, schema);
 
     fbb.finish(root, None);
@@ -60,14 +110,18 @@ pub fn metadata_to_fb<'a>(
     fbb.create_vector(&custom_metadata)
 }
 
+#[deprecated(
+    since = "54.0.0",
+    note = "Use `IpcSchemaConverter` instead. This function will be removed in the next release."
+)]
 pub fn schema_to_fb_offset<'a>(
     fbb: &mut FlatBufferBuilder<'a>,
     schema: &Schema,
 ) -> WIPOffset<crate::Schema<'a>> {
     let fields = schema
         .fields()
         .iter()
-        .map(|field| build_field(fbb, field))
+        .map(|field| build_field(fbb, &mut None, field))
         .collect::<Vec<_>>();
     let fb_field_list = fbb.create_vector(&fields);
 
@@ -424,6 +478,7 @@ pub(crate) struct FBFieldType<'b> {
 /// Create an IPC Field from an Arrow Field
 pub(crate) fn build_field<'a>(
     fbb: &mut FlatBufferBuilder<'a>,
+    dictionary_tracker: &mut Option<&mut DictionaryTracker>,
     field: &Field,
 ) -> WIPOffset<crate::Field<'a>> {
     // Optional custom metadata.
@@ -433,19 +488,29 @@ pub(crate) fn build_field<'a>(
     };
 
     let fb_field_name = fbb.create_string(field.name().as_str());
-    let field_type = get_fb_field_type(field.data_type(), fbb);
+    let field_type = get_fb_field_type(field.data_type(), dictionary_tracker, fbb);
 
     let fb_dictionary = if let Dictionary(index_type, _) = field.data_type() {
-        Some(get_fb_dictionary(
-            index_type,
-            field
-                .dict_id()
-                .expect("All Dictionary types have `dict_id`"),
-            field
-                .dict_is_ordered()
-                .expect("All Dictionary types have `dict_is_ordered`"),
-            fbb,
-        ))
+        match dictionary_tracker {
+            Some(tracker) => Some(get_fb_dictionary(
+                index_type,
+                tracker.set_dict_id(field),
+                field
+                    .dict_is_ordered()
+                    .expect("All Dictionary types have `dict_is_ordered`"),
+                fbb,
+            )),
+            None => Some(get_fb_dictionary(
+                index_type,
+                field
+                    .dict_id()
+                    .expect("Dictionary type must have a dictionary id"),
+                field
+                    .dict_is_ordered()
+                    .expect("All Dictionary types have `dict_is_ordered`"),
+                fbb,
+            )),
+        }
     } else {
         None
     };
@@ -473,6 +538,7 @@ pub(crate) fn build_field<'a>(
 /// Get the IPC type of a data type
 pub(crate) fn get_fb_field_type<'a>(
     data_type: &DataType,
+    dictionary_tracker: &mut Option<&mut DictionaryTracker>,
     fbb: &mut FlatBufferBuilder<'a>,
 ) -> FBFieldType<'a> {
     // some IPC implementations expect an empty list for child data, instead of a null value.
@@ -673,7 +739,7 @@ pub(crate) fn get_fb_field_type<'a>(
             }
         }
         List(ref list_type) => {
-            let child = build_field(fbb, list_type);
+            let child = build_field(fbb, dictionary_tracker, list_type);
             FBFieldType {
                 type_type: crate::Type::List,
                 type_: crate::ListBuilder::new(fbb).finish().as_union_value(),
@@ -682,15 +748,15 @@ pub(crate) fn get_fb_field_type<'a>(
         }
         ListView(_) | LargeListView(_) => unimplemented!("ListView/LargeListView not implemented"),
         LargeList(ref list_type) => {
-            let child = build_field(fbb, list_type);
+            let child = build_field(fbb, dictionary_tracker, list_type);
             FBFieldType {
                 type_type: crate::Type::LargeList,
                 type_: crate::LargeListBuilder::new(fbb).finish().as_union_value(),
                 children: Some(fbb.create_vector(&[child])),
             }
         }
         FixedSizeList(ref list_type, len) => {
-            let child = build_field(fbb, list_type);
+            let child = build_field(fbb, dictionary_tracker, list_type);
             let mut builder = crate::FixedSizeListBuilder::new(fbb);
             builder.add_listSize(*len);
             FBFieldType {
@@ -703,7 +769,7 @@ pub(crate) fn get_fb_field_type<'a>(
             // struct's fields are children
             let mut children = vec![];
             for field in fields {
-                children.push(build_field(fbb, field));
+                children.push(build_field(fbb, dictionary_tracker, field));
             }
             FBFieldType {
                 type_type: crate::Type::Struct_,
@@ -712,8 +778,8 @@ pub(crate) fn get_fb_field_type<'a>(
             }
         }
         RunEndEncoded(run_ends, values) => {
-            let run_ends_field = build_field(fbb, run_ends);
-            let values_field = build_field(fbb, values);
+            let run_ends_field = build_field(fbb, dictionary_tracker, run_ends);
+            let values_field = build_field(fbb, dictionary_tracker, values);
             let children = [run_ends_field, values_field];
             FBFieldType {
                 type_type: crate::Type::RunEndEncoded,
@@ -724,7 +790,7 @@ pub(crate) fn get_fb_field_type<'a>(
             }
         }
         Map(map_field, keys_sorted) => {
-            let child = build_field(fbb, map_field);
+            let child = build_field(fbb, dictionary_tracker, map_field);
             let mut field_type = crate::MapBuilder::new(fbb);
             field_type.add_keysSorted(*keys_sorted);
             FBFieldType {
@@ -737,7 +803,7 @@ pub(crate) fn get_fb_field_type<'a>(
             // In this library, the dictionary "type" is a logical construct. Here we
             // pass through to the value type, as we've already captured the index
             // type in the DictionaryEncoding metadata in the parent field
-            get_fb_field_type(value_type, fbb)
+            get_fb_field_type(value_type, dictionary_tracker, fbb)
         }
         Decimal128(precision, scale) => {
             let mut builder = crate::DecimalBuilder::new(fbb);
@@ -764,7 +830,7 @@ pub(crate) fn get_fb_field_type<'a>(
         Union(fields, mode) => {
             let mut children = vec![];
             for (_, field) in fields.iter() {
-                children.push(build_field(fbb, field));
+                children.push(build_field(fbb, dictionary_tracker, field));
             }
 
             let union_mode = match mode {
@@ -1067,7 +1133,9 @@ mod tests {
             md,
         );
 
-        let fb = schema_to_fb(&schema);
+        let mut dictionary_tracker = DictionaryTracker::new(true);
+        let mut converter = IpcSchemaConverter::new(&mut dictionary_tracker);
+        let fb = converter.schema_to_fb(&schema);
 
         // read back fields
         let ipc = crate::root_as_schema(fb.finished_data()).unwrap();
@@ -1098,9 +1166,14 @@ mod tests {
 
         // generate same message with Rust
         let data_gen = crate::writer::IpcDataGenerator::default();
+        let mut dictionary_tracker = DictionaryTracker::new(true);
         let arrow_schema = Schema::new(vec![Field::new("field1", DataType::UInt32, false)]);
         let bytes = data_gen
-            .schema_to_bytes(&arrow_schema, &crate::writer::IpcWriteOptions::default())
+            .schema_to_bytes_with_dictionary_tracker(
+                &arrow_schema,
+                &mut dictionary_tracker,
+                &crate::writer::IpcWriteOptions::default(),
+            )
             .ipc_message;
 
         let ipc2 = crate::root_as_message(&bytes).unwrap();

arrow-ipc/src/reader.rs

@@ -2283,4 +2283,59 @@ mod tests {
         let err = reader.next().unwrap().unwrap_err();
         assert!(matches!(err, ArrowError::InvalidArgumentError(_)));
     }
+
+    #[test]
+    fn test_same_dict_id_without_preserve() {
+        let batch = RecordBatch::try_new(
+            Arc::new(Schema::new(
+                ["a", "b"]
+                    .iter()
+                    .map(|name| {
+                        Field::new_dict(
+                            name.to_string(),
+                            DataType::Dictionary(
+                                Box::new(DataType::Int32),
+                                Box::new(DataType::Utf8),
+                            ),
+                            true,
+                            0,
+                            false,
+                        )
+                    })
+                    .collect::<Vec<Field>>(),
+            )),
+            vec![
+                Arc::new(
+                    vec![Some("c"), Some("d")]
+                        .into_iter()
+                        .collect::<DictionaryArray<Int32Type>>(),
+                ) as ArrayRef,
+                Arc::new(
+                    vec![Some("e"), Some("f")]
+                        .into_iter()
+                        .collect::<DictionaryArray<Int32Type>>(),
+                ) as ArrayRef,
+            ],
+        )
+        .expect("Failed to create RecordBatch");
+
+        // serialize the record batch as an IPC stream
+        let mut buf = vec![];
+        {
+            let mut writer = crate::writer::StreamWriter::try_new_with_options(
+                &mut buf,
+                batch.schema().as_ref(),
+                crate::writer::IpcWriteOptions::default().with_preserve_dict_id(false),
+            )
+            .expect("Failed to create StreamWriter");
+            writer.write(&batch).expect("Failed to write RecordBatch");
+            writer.finish().expect("Failed to finish StreamWriter");
+        }
+
+        StreamReader::try_new(std::io::Cursor::new(buf), None)
+            .expect("Failed to create StreamReader")
+            .for_each(|decoded_batch| {
+                assert_eq!(decoded_batch.expect("Failed to read RecordBatch"), batch);
+            });
+    }
 }