From 4655a3cfe3b49d745413394e2c8426216ec54b5a Mon Sep 17 00:00:00 2001
From: chielP <chielpeters@live.nl>
Date: Sun, 14 Apr 2024 19:21:28 +0200
Subject: [PATCH] rebase

---
 crates/polars-arrow/Cargo.toml                       |  2 +-
 crates/polars-arrow/src/io/ipc/read/file.rs          |  3 ---
 crates/polars-arrow/src/io/ipc/read/flight_reader.rs | 10 ++++++----
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/crates/polars-arrow/Cargo.toml b/crates/polars-arrow/Cargo.toml
index 529a7b7003c8..efce51cdf4c3 100644
--- a/crates/polars-arrow/Cargo.toml
+++ b/crates/polars-arrow/Cargo.toml
@@ -120,7 +120,7 @@ io_ipc = ["arrow-format", "polars-error/arrow-format"]
 io_ipc_write_async = ["io_ipc", "futures"]
 io_ipc_read_async = ["io_ipc", "futures", "async-stream"]
 io_ipc_compression = ["lz4", "zstd", "io_ipc"]
-io_flight = ["io_ipc", "arrow-format/flight-data","io_ipc_read_async"]
+io_flight = ["io_ipc", "arrow-format/flight-data", "io_ipc_read_async"]
 
 io_avro = ["avro-schema", "polars-error/avro-schema"]
 io_avro_compression = [
diff --git a/crates/polars-arrow/src/io/ipc/read/file.rs b/crates/polars-arrow/src/io/ipc/read/file.rs
index dc654af214e6..80e8c7ef703a 100644
--- a/crates/polars-arrow/src/io/ipc/read/file.rs
+++ b/crates/polars-arrow/src/io/ipc/read/file.rs
@@ -75,9 +75,6 @@ fn read_dictionary_block<R: Read + Seek>(
     message_scratch: &mut Vec<u8>,
     dictionary_scratch: &mut Vec<u8>,
 ) -> PolarsResult<()> {
-    let message = read_ipc_message_from_block(reader, block, message_scratch)?;
-    let batch = get_dictionary_batch(&message)?;
-
     let offset: u64 = block
         .offset
         .try_into()
diff --git a/crates/polars-arrow/src/io/ipc/read/flight_reader.rs b/crates/polars-arrow/src/io/ipc/read/flight_reader.rs
index 99a8416ef7f1..a8ec5a0a6877 100644
--- a/crates/polars-arrow/src/io/ipc/read/flight_reader.rs
+++ b/crates/polars-arrow/src/io/ipc/read/flight_reader.rs
@@ -5,12 +5,11 @@ use polars_error::{polars_bail, polars_err, PolarsResult};
 
 use super::OutOfSpecKind;
 use crate::array::Array;
-use crate::chunk::Chunk;
 use crate::datatypes::ArrowSchema;
 use crate::io::ipc::read::schema::deserialize_stream_metadata;
 use crate::io::ipc::read::{read_dictionary, read_record_batch, Dictionaries, StreamMetadata};
 use crate::io::ipc::write::common::EncodedData;
-
+use crate::record_batch::RecordBatch;
 pub struct FlightStreamReader {
     metadata: Option<StreamMetadata>,
     dictionaries: Dictionaries,
@@ -33,7 +32,10 @@ impl FlightStreamReader {
             .map(|metadata| &metadata.schema)
     }
 
-    pub fn parse(&mut self, data: EncodedData) -> PolarsResult<Option<Chunk<Box<dyn Array>>>> {
+    pub fn parse(
+        &mut self,
+        data: EncodedData,
+    ) -> PolarsResult<Option<RecordBatch<Box<dyn Array>>>> {
         // First message should be the schema
         if self.metadata.is_none() {
             let message = arrow_format::ipc::MessageRef::read_as_root(data.ipc_message.as_ref())
@@ -63,7 +65,7 @@ fn convert_to_arrow_chunk(
     metadata: &StreamMetadata,
     dictionaries: &mut Dictionaries,
     flight_data: EncodedData,
-) -> PolarsResult<Option<Chunk<Box<dyn Array>>>> {
+) -> PolarsResult<Option<RecordBatch<Box<dyn Array>>>> {
     let EncodedData {
         ipc_message,
         arrow_data,