laminar_connectors/files/arrow_ipc_codec/
mod.rs

1//! Arrow IPC file format decoder and encoder.
2
3use std::io::Cursor;
4
5use arrow_array::RecordBatch;
6use arrow_schema::SchemaRef;
7
8use crate::schema::error::{SchemaError, SchemaResult};
9use crate::schema::traits::{FormatDecoder, FormatEncoder};
10use crate::schema::types::RawRecord;
11
12/// Decodes Arrow IPC file bytes into `RecordBatch`es.
13///
14/// The constructor schema is used for `output_schema()` and empty-batch
15/// returns. Actual decoded batches carry the file's embedded schema
16/// (same contract as `ParquetDecoder`).
17pub struct ArrowIpcDecoder {
18    schema: SchemaRef,
19}
20
21impl ArrowIpcDecoder {
22    /// Creates a decoder with the given declared schema.
23    #[must_use]
24    pub fn new(schema: SchemaRef) -> Self {
25        Self { schema }
26    }
27}
28
29impl std::fmt::Debug for ArrowIpcDecoder {
30    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        f.debug_struct("ArrowIpcDecoder")
32            .field("schema", &self.schema)
33            .finish()
34    }
35}
36
37impl FormatDecoder for ArrowIpcDecoder {
38    fn output_schema(&self) -> SchemaRef {
39        self.schema.clone()
40    }
41
42    fn decode_batch(&self, records: &[RawRecord]) -> SchemaResult<RecordBatch> {
43        if records.is_empty() {
44            return Ok(RecordBatch::new_empty(self.schema.clone()));
45        }
46
47        let mut combined = Vec::with_capacity(records.iter().map(|r| r.value.len()).sum());
48        for record in records {
49            combined.extend_from_slice(&record.value);
50        }
51
52        let cursor = Cursor::new(&combined);
53        let reader = arrow_ipc::reader::FileReader::try_new(cursor, None)
54            .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC read error: {e}")))?;
55
56        let file_schema = reader.schema();
57
58        let mut batches = Vec::new();
59        for batch_result in reader {
60            let batch = batch_result
61                .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC batch error: {e}")))?;
62            batches.push(batch);
63        }
64
65        if batches.is_empty() {
66            return Ok(RecordBatch::new_empty(file_schema));
67        }
68
69        if batches.len() == 1 {
70            return Ok(batches.into_iter().next().unwrap());
71        }
72
73        arrow_select::concat::concat_batches(&file_schema, &batches)
74            .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC concat error: {e}")))
75    }
76
77    fn format_name(&self) -> &'static str {
78        "arrow_ipc"
79    }
80}
81
82/// Encodes `RecordBatch`es into Arrow IPC file format bytes.
83#[derive(Debug)]
84pub struct ArrowIpcEncoder {
85    schema: SchemaRef,
86}
87
88impl ArrowIpcEncoder {
89    /// Creates a new Arrow IPC encoder for the given schema.
90    #[must_use]
91    pub fn new(schema: SchemaRef) -> Self {
92        Self { schema }
93    }
94}
95
96impl FormatEncoder for ArrowIpcEncoder {
97    fn input_schema(&self) -> SchemaRef {
98        self.schema.clone()
99    }
100
101    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>> {
102        if batch.num_rows() == 0 {
103            return Ok(Vec::new());
104        }
105
106        let mut buf = Vec::new();
107        {
108            let mut writer = arrow_ipc::writer::FileWriter::try_new(&mut buf, &batch.schema())
109                .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC writer init: {e}")))?;
110            writer
111                .write(batch)
112                .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC write error: {e}")))?;
113            writer
114                .finish()
115                .map_err(|e| SchemaError::DecodeError(format!("Arrow IPC finish error: {e}")))?;
116        }
117
118        Ok(vec![buf])
119    }
120
121    fn format_name(&self) -> &'static str {
122        "arrow_ipc"
123    }
124}
125#[cfg(test)]
126mod tests;
laminar_connectors/files/arrow_ipc_codec/mod.rs

laminar_connectors/files/arrow_ipc_codec/
mod.rs
