Upgrade datafusion (#2056)

nautechsystems · Nov 17, 2024 · 3896902 · 3896902
1 parent 762d723
commit 3896902
Show file tree

Hide file tree

Showing 3 changed files with 80 additions and 47 deletions.
diff --git a/nautilus_core/Cargo.lock b/nautilus_core/Cargo.lock
diff --git a/nautilus_core/persistence/Cargo.toml b/nautilus_core/persistence/Cargo.toml
@@ -24,7 +24,7 @@ tokio = { workspace = true }
 thiserror = { workspace = true }
 binary-heap-plus = "0.5.0"
 compare = "0.1.0"
-datafusion = { version = "42.2.0", default-features = false, features = ["compression", "regex_expressions", "unicode_expressions", "pyarrow"] }
+datafusion = { version = "43.0.0", default-features = false, features = ["compression", "regex_expressions", "unicode_expressions", "pyarrow"] }
 
 [dev-dependencies]
 nautilus-test-kit = { path = "../test_kit" }

diff --git a/nautilus_core/serialization/src/arrow/trade.rs b/nautilus_core/serialization/src/arrow/trade.rs
@@ -16,7 +16,7 @@
 use std::{collections::HashMap, str::FromStr, sync::Arc};
 
 use arrow::{
-    array::{Int64Array, StringArray, StringBuilder, UInt64Array, UInt8Array},
+    array::{Int64Array, StringArray, StringBuilder, StringViewArray, UInt64Array, UInt8Array},
     datatypes::{DataType, Field, Schema},
     error::ArrowError,
     record_batch::RecordBatch,
@@ -130,10 +130,27 @@ impl DecodeFromRecordBatch for TradeTick {
         let size_values = extract_column::<UInt64Array>(cols, "size", 1, DataType::UInt64)?;
         let aggressor_side_values =
             extract_column::<UInt8Array>(cols, "aggressor_side", 2, DataType::UInt8)?;
-        let trade_id_values = extract_column::<StringArray>(cols, "trade_id", 3, DataType::Utf8)?;
         let ts_event_values = extract_column::<UInt64Array>(cols, "ts_event", 4, DataType::UInt64)?;
         let ts_init_values = extract_column::<UInt64Array>(cols, "ts_init", 5, DataType::UInt64)?;
 
+        // Datafusion reads trade_ids as StringView
+        let trade_id_values: Vec<TradeId> = if record_batch
+            .schema()
+            .field_with_name("trade_id")?
+            .data_type()
+            == &DataType::Utf8View
+        {
+            extract_column::<StringViewArray>(cols, "trade_id", 3, DataType::Utf8View)?
+                .iter()
+                .map(|id| TradeId::from(id.unwrap()))
+                .collect()
+        } else {
+            extract_column::<StringArray>(cols, "trade_id", 3, DataType::Utf8)?
+                .iter()
+                .map(|id| TradeId::from(id.unwrap()))
+                .collect()
+        };
+
         let result: Result<Vec<Self>, EncodingError> = (0..record_batch.num_rows())
             .map(|i| {
                 let price = Price::from_raw(price_values.value(i), price_precision);
@@ -146,7 +163,7 @@ impl DecodeFromRecordBatch for TradeTick {
                             format!("Invalid enum value, was {aggressor_side_value}"),
                         )
                     })?;
-                let trade_id = TradeId::from(trade_id_values.value(i));
+                let trade_id = trade_id_values[i];
                 let ts_event = ts_event_values.value(i).into();
                 let ts_init = ts_init_values.value(i).into();
 
@@ -184,7 +201,7 @@ mod tests {
     use std::sync::Arc;
 
     use arrow::{
-        array::{Array, Int64Array, StringArray, UInt64Array, UInt8Array},
+        array::{Array, Int64Array, UInt64Array, UInt8Array},
         record_batch::RecordBatch,
     };
     use rstest::rstest;