Make dict ID only an IPC concern

Author: brancz

arrow-flight/src/decode.rs

@@ -276,6 +276,7 @@ impl FlightDataDecoder {
 
                 self.state = Some(FlightStreamState {
                     schema: Arc::clone(&schema),
+                    schema_message: data.clone(),
                     dictionaries_by_field,
                 });
                 Ok(Some(DecodedFlightData::new_schema(data, schema)))
@@ -296,10 +297,15 @@ impl FlightDataDecoder {
                     )
                 })?;
 
+                let ipc_schema = arrow_ipc::root_as_message(&state.schema_message.data_header)
+                    .unwrap()
+                    .header_as_schema()
+                    .unwrap();
+
                 arrow_ipc::reader::read_dictionary(
                     &buffer,
                     dictionary_batch,
-                    &state.schema,
+                    ipc_schema,
                     &mut state.dictionaries_by_field,
                     &message.version(),
                 )
@@ -319,8 +325,14 @@ impl FlightDataDecoder {
                     ));
                 };
 
+                let ipc_schema = arrow_ipc::root_as_message(&state.schema_message.data_header)
+                    .unwrap()
+                    .header_as_schema()
+                    .unwrap();
+
                 let batch = flight_data_to_arrow_batch(
                     &data,
+                    ipc_schema,
                     Arc::clone(&state.schema),
                     &state.dictionaries_by_field,
                 )
@@ -376,6 +388,7 @@ impl futures::Stream for FlightDataDecoder {
 #[derive(Debug)]
 struct FlightStreamState {
     schema: SchemaRef,
+    schema_message: FlightData,
     dictionaries_by_field: HashMap<i64, ArrayRef>,
 }
 

arrow-flight/src/encode.rs

@@ -535,15 +535,13 @@ fn prepare_field_for_flight(
                 )
                 .with_metadata(field.metadata().clone())
             } else {
-                #[allow(deprecated)]
-                let dict_id = dictionary_tracker.set_dict_id(field.as_ref());
-
+                dictionary_tracker.next_dict_id();
                 #[allow(deprecated)]
                 Field::new_dict(
                     field.name(),
                     field.data_type().clone(),
                     field.is_nullable(),
-                    dict_id,
+                    0,
                     field.dict_is_ordered().unwrap_or_default(),
                 )
                 .with_metadata(field.metadata().clone())
@@ -585,14 +583,13 @@ fn prepare_schema_for_flight(
                     )
                     .with_metadata(field.metadata().clone())
                 } else {
-                    #[allow(deprecated)]
-                    let dict_id = dictionary_tracker.set_dict_id(field.as_ref());
+                    dictionary_tracker.next_dict_id();
                     #[allow(deprecated)]
                     Field::new_dict(
                         field.name(),
                         field.data_type().clone(),
                         field.is_nullable(),
-                        dict_id,
+                        0,
                         field.dict_is_ordered().unwrap_or_default(),
                     )
                     .with_metadata(field.metadata().clone())
@@ -654,16 +651,10 @@ struct FlightIpcEncoder {
 
 impl FlightIpcEncoder {
     fn new(options: IpcWriteOptions, error_on_replacement: bool) -> Self {
-        #[allow(deprecated)]
-        let preserve_dict_id = options.preserve_dict_id();
         Self {
             options,
             data_gen: IpcDataGenerator::default(),
-            #[allow(deprecated)]
-            dictionary_tracker: DictionaryTracker::new_with_preserve_dict_id(
-                error_on_replacement,
-                preserve_dict_id,
-            ),
+            dictionary_tracker: DictionaryTracker::new(error_on_replacement),
         }
     }
 
@@ -1547,9 +1538,8 @@ mod tests {
     async fn verify_flight_round_trip(mut batches: Vec<RecordBatch>) {
         let expected_schema = batches.first().unwrap().schema();
 
-        #[allow(deprecated)]
         let encoder = FlightDataEncoderBuilder::default()
-            .with_options(IpcWriteOptions::default().with_preserve_dict_id(false))
+            .with_options(IpcWriteOptions::default())
             .with_dictionary_handling(DictionaryHandling::Resend)
             .build(futures::stream::iter(batches.clone().into_iter().map(Ok)));
 
@@ -1575,8 +1565,7 @@ mod tests {
             HashMap::from([("some_key".to_owned(), "some_value".to_owned())]),
         );
 
-        #[allow(deprecated)]
-        let mut dictionary_tracker = DictionaryTracker::new_with_preserve_dict_id(false, true);
+        let mut dictionary_tracker = DictionaryTracker::new(false);
 
         let got = prepare_schema_for_flight(&schema, &mut dictionary_tracker, false);
         assert!(got.metadata().contains_key("some_key"));
@@ -1606,9 +1595,7 @@ mod tests {
         options: &IpcWriteOptions,
     ) -> (Vec<FlightData>, FlightData) {
         let data_gen = IpcDataGenerator::default();
-        #[allow(deprecated)]
-        let mut dictionary_tracker =
-            DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());
+        let mut dictionary_tracker = DictionaryTracker::new(false);
 
         let (encoded_dictionaries, encoded_batch) = data_gen
             .encoded_batch(batch, &mut dictionary_tracker, options)

arrow-flight/src/lib.rs

@@ -149,9 +149,7 @@ pub struct IpcMessage(pub Bytes);
 
 fn flight_schema_as_encoded_data(arrow_schema: &Schema, options: &IpcWriteOptions) -> EncodedData {
     let data_gen = writer::IpcDataGenerator::default();
-    #[allow(deprecated)]
-    let mut dict_tracker =
-        writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());
+    let mut dict_tracker = writer::DictionaryTracker::new(false);
     data_gen.schema_to_bytes_with_dictionary_tracker(arrow_schema, &mut dict_tracker, options)
 }
 

arrow-flight/src/sql/client.rs

@@ -707,6 +707,7 @@ pub enum ArrowFlightData {
 pub fn arrow_data_from_flight_data(
     flight_data: FlightData,
     arrow_schema_ref: &SchemaRef,
+    ipc_schema: arrow_ipc::Schema,
 ) -> Result<ArrowFlightData, ArrowError> {
     let ipc_message = root_as_message(&flight_data.data_header[..])
         .map_err(|err| ArrowError::ParseError(format!("Unable to get root as message: {err:?}")))?;
@@ -723,6 +724,7 @@ pub fn arrow_data_from_flight_data(
             let record_batch = read_record_batch(
                 &Buffer::from(flight_data.data_body),
                 ipc_record_batch,
+                ipc_schema,
                 arrow_schema_ref.clone(),
                 &dictionaries_by_field,
                 None,

arrow-flight/src/utils.rs

@@ -44,7 +44,8 @@ pub fn flight_data_to_batches(flight_data: &[FlightData]) -> Result<Vec<RecordBa
     let mut batches = vec![];
     let dictionaries_by_id = HashMap::new();
     for datum in flight_data[1..].iter() {
-        let batch = flight_data_to_arrow_batch(datum, schema.clone(), &dictionaries_by_id)?;
+        let batch =
+            flight_data_to_arrow_batch(datum, ipc_schema, schema.clone(), &dictionaries_by_id)?;
         batches.push(batch);
     }
     Ok(batches)
@@ -53,6 +54,7 @@ pub fn flight_data_to_batches(flight_data: &[FlightData]) -> Result<Vec<RecordBa
 /// Convert `FlightData` (with supplied schema and dictionaries) to an arrow `RecordBatch`.
 pub fn flight_data_to_arrow_batch(
     data: &FlightData,
+    ipc_schema: arrow_ipc::Schema,
     schema: SchemaRef,
     dictionaries_by_id: &HashMap<i64, ArrayRef>,
 ) -> Result<RecordBatch, ArrowError> {
@@ -71,6 +73,7 @@ pub fn flight_data_to_arrow_batch(
             reader::read_record_batch(
                 &Buffer::from(data.data_body.as_ref()),
                 batch,
+                ipc_schema,
                 schema,
                 dictionaries_by_id,
                 None,
@@ -90,9 +93,7 @@ pub fn batches_to_flight_data(
     let mut flight_data = vec![];
 
     let data_gen = writer::IpcDataGenerator::default();
-    #[allow(deprecated)]
-    let mut dictionary_tracker =
-        writer::DictionaryTracker::new_with_preserve_dict_id(false, options.preserve_dict_id());
+    let mut dictionary_tracker = writer::DictionaryTracker::new(false);
 
     for batch in batches.iter() {
         let (encoded_dictionaries, encoded_batch) =

arrow-integration-test/Cargo.toml

@@ -38,6 +38,7 @@ all-features = true
 [dependencies]
 arrow = { workspace = true }
 arrow-buffer = { workspace = true }
+arrow-ipc = { workspace = true }
 hex = { version = "0.4", default-features = false, features = ["std"] }
 serde = { version = "1.0", default-features = false, features = ["rc", "derive"] }
 serde_json = { version = "1.0", default-features = false, features = ["std"] }

arrow-integration-test/src/field.rs

@@ -18,6 +18,7 @@
 use crate::{data_type_from_json, data_type_to_json};
 use arrow::datatypes::{DataType, Field};
 use arrow::error::{ArrowError, Result};
+use arrow_ipc::writer::DictionaryTracker;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -218,7 +219,6 @@ pub fn field_from_json(json: &serde_json::Value) -> Result<Field> {
                 _ => data_type,
             };
 
-            let mut dict_id = 0;
             let mut dict_is_ordered = false;
 
             let data_type = match map.get("dictionary") {
@@ -231,14 +231,6 @@ pub fn field_from_json(json: &serde_json::Value) -> Result<Field> {
                             ));
                         }
                     };
-                    dict_id = match dictionary.get("id") {
-                        Some(Value::Number(n)) => n.as_i64().unwrap(),
-                        _ => {
-                            return Err(ArrowError::ParseError(
-                                "Field missing 'id' attribute".to_string(),
-                            ));
-                        }
-                    };
                     dict_is_ordered = match dictionary.get("isOrdered") {
                         Some(&Value::Bool(n)) => n,
                         _ => {
@@ -253,7 +245,7 @@ pub fn field_from_json(json: &serde_json::Value) -> Result<Field> {
             };
 
             #[allow(deprecated)]
-            let mut field = Field::new_dict(name, data_type, nullable, dict_id, dict_is_ordered);
+            let mut field = Field::new_dict(name, data_type, nullable, 0, dict_is_ordered);
             field.set_metadata(metadata);
             Ok(field)
         }
@@ -264,27 +256,28 @@ pub fn field_from_json(json: &serde_json::Value) -> Result<Field> {
 }
 
 /// Generate a JSON representation of the `Field`.
-pub fn field_to_json(field: &Field) -> serde_json::Value {
+pub fn field_to_json(dict_tracker: &mut DictionaryTracker, field: &Field) -> serde_json::Value {
     let children: Vec<serde_json::Value> = match field.data_type() {
-        DataType::Struct(fields) => fields.iter().map(|x| field_to_json(x.as_ref())).collect(),
+        DataType::Struct(fields) => fields
+            .iter()
+            .map(|x| field_to_json(dict_tracker, x.as_ref()))
+            .collect(),
         DataType::List(field)
         | DataType::LargeList(field)
         | DataType::FixedSizeList(field, _)
-        | DataType::Map(field, _) => vec![field_to_json(field)],
+        | DataType::Map(field, _) => vec![field_to_json(dict_tracker, field)],
         _ => vec![],
     };
 
     match field.data_type() {
         DataType::Dictionary(ref index_type, ref value_type) => {
-            #[allow(deprecated)]
-            let dict_id = field.dict_id().unwrap();
             serde_json::json!({
                 "name": field.name(),
                 "nullable": field.is_nullable(),
                 "type": data_type_to_json(value_type),
                 "children": children,
                 "dictionary": {
-                    "id": dict_id,
+                    "id": dict_tracker.next_dict_id(),
                     "indexType": data_type_to_json(index_type),
                     "isOrdered": field.dict_is_ordered().unwrap(),
                 }
@@ -345,7 +338,8 @@ mod tests {
             }"#,
         )
         .unwrap();
-        assert_eq!(value, field_to_json(&f));
+        let mut dictionary_tracker = DictionaryTracker::new(false);
+        assert_eq!(value, field_to_json(&mut dictionary_tracker, &f));
     }
 
     #[test]
@@ -398,7 +392,8 @@ mod tests {
             }"#,
         )
         .unwrap();
-        assert_eq!(value, field_to_json(&f));
+        let mut dictionary_tracker = DictionaryTracker::new(false);
+        assert_eq!(value, field_to_json(&mut dictionary_tracker, &f));
     }
 
     #[test]
@@ -415,7 +410,8 @@ mod tests {
             }"#,
         )
         .unwrap();
-        assert_eq!(value, field_to_json(&f));
+        let mut dictionary_tracker = DictionaryTracker::new(false);
+        assert_eq!(value, field_to_json(&mut dictionary_tracker, &f));
     }
     #[test]
     fn parse_struct_from_json() {