ipfrs_core/
arrow.rs

1//! Apache Arrow memory layout integration for zero-copy tensor access.
2//!
3//! This module provides conversions between IPFRS tensor types and Apache Arrow arrays,
4//! enabling zero-copy interoperability with the Arrow ecosystem (Parquet, Flight, etc.).
5//!
6//! ## Example
7//!
8//! ```rust
9//! use ipfrs_core::arrow::{TensorBlockArrowExt, arrow_to_tensor_block};
10//! use ipfrs_core::tensor::{TensorBlock, TensorDtype, TensorShape};
11//! use bytes::Bytes;
12//! use arrow_array::Float32Array;
13//!
14//! // Convert Arrow array to TensorBlock (zero-copy)
15//! let arrow_array = Float32Array::from(vec![1.0f32, 2.0, 3.0, 4.0]);
16//! let tensor = arrow_to_tensor_block(&arrow_array, TensorShape::new(vec![2, 2])).unwrap();
17//!
18//! // Convert TensorBlock back to Arrow array
19//! let arrow_back = tensor.to_arrow_array().unwrap();
20//! ```
21
22use crate::error::{Error, Result};
23use crate::tensor::{TensorBlock, TensorDtype, TensorShape};
24use arrow_array::{
25    Array, ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, Int8Array,
26    UInt32Array, UInt8Array,
27};
28use arrow_buffer::Buffer;
29use arrow_schema::{DataType, Field, Schema};
30use bytes::Bytes;
31use std::sync::Arc;
32
33/// Extension trait for TensorBlock to provide Arrow conversions
34pub trait TensorBlockArrowExt {
35    /// Convert to an Arrow array (zero-copy when possible)
36    fn to_arrow_array(&self) -> Result<ArrayRef>;
37
38    /// Convert to an Arrow Field (for schema)
39    fn to_arrow_field(&self, name: &str) -> Field;
40
41    /// Convert to an Arrow Schema
42    fn to_arrow_schema(&self, field_name: &str) -> Schema;
43}
44
45impl TensorBlockArrowExt for TensorBlock {
46    fn to_arrow_array(&self) -> Result<ArrayRef> {
47        let metadata = self.metadata();
48        let data = self.data();
49
50        match metadata.dtype {
51            TensorDtype::F32 => {
52                let buffer = Buffer::from(data.clone());
53                let array = Float32Array::new(buffer.into(), None);
54                Ok(Arc::new(array) as ArrayRef)
55            }
56            TensorDtype::F64 => {
57                let buffer = Buffer::from(data.clone());
58                let array = Float64Array::new(buffer.into(), None);
59                Ok(Arc::new(array) as ArrayRef)
60            }
61            TensorDtype::I8 => {
62                let buffer = Buffer::from(data.clone());
63                let array = Int8Array::new(buffer.into(), None);
64                Ok(Arc::new(array) as ArrayRef)
65            }
66            TensorDtype::I32 => {
67                let buffer = Buffer::from(data.clone());
68                let array = Int32Array::new(buffer.into(), None);
69                Ok(Arc::new(array) as ArrayRef)
70            }
71            TensorDtype::I64 => {
72                let buffer = Buffer::from(data.clone());
73                let array = Int64Array::new(buffer.into(), None);
74                Ok(Arc::new(array) as ArrayRef)
75            }
76            TensorDtype::U8 => {
77                let buffer = Buffer::from(data.clone());
78                let array = UInt8Array::new(buffer.into(), None);
79                Ok(Arc::new(array) as ArrayRef)
80            }
81            TensorDtype::U32 => {
82                let buffer = Buffer::from(data.clone());
83                let array = UInt32Array::new(buffer.into(), None);
84                Ok(Arc::new(array) as ArrayRef)
85            }
86            TensorDtype::Bool => {
87                // Boolean arrays are stored as bit-packed in Arrow
88                let bytes: Vec<u8> = data.to_vec();
89                let array = BooleanArray::from(bytes.iter().map(|&b| b != 0).collect::<Vec<_>>());
90                Ok(Arc::new(array) as ArrayRef)
91            }
92            TensorDtype::F16 => {
93                // Arrow doesn't have native F16 support, convert to F32
94                Err(Error::InvalidInput(
95                    "F16 not directly supported by Arrow, use F32 instead".to_string(),
96                ))
97            }
98        }
99    }
100
101    fn to_arrow_field(&self, name: &str) -> Field {
102        let metadata = self.metadata();
103        let arrow_dtype = tensor_dtype_to_arrow(&metadata.dtype);
104        Field::new(name, arrow_dtype, false)
105    }
106
107    fn to_arrow_schema(&self, field_name: &str) -> Schema {
108        Schema::new(vec![self.to_arrow_field(field_name)])
109    }
110}
111
112/// Convert Arrow DataType to TensorDtype
113pub fn arrow_dtype_to_tensor(dtype: &DataType) -> Result<TensorDtype> {
114    match dtype {
115        DataType::Float32 => Ok(TensorDtype::F32),
116        DataType::Float64 => Ok(TensorDtype::F64),
117        DataType::Int8 => Ok(TensorDtype::I8),
118        DataType::Int32 => Ok(TensorDtype::I32),
119        DataType::Int64 => Ok(TensorDtype::I64),
120        DataType::UInt8 => Ok(TensorDtype::U8),
121        DataType::UInt32 => Ok(TensorDtype::U32),
122        DataType::Boolean => Ok(TensorDtype::Bool),
123        _ => Err(Error::InvalidInput(format!(
124            "Unsupported Arrow dtype: {:?}",
125            dtype
126        ))),
127    }
128}
129
130/// Convert TensorDtype to Arrow DataType
131pub fn tensor_dtype_to_arrow(dtype: &TensorDtype) -> DataType {
132    match dtype {
133        TensorDtype::F32 => DataType::Float32,
134        TensorDtype::F64 => DataType::Float64,
135        TensorDtype::I8 => DataType::Int8,
136        TensorDtype::I32 => DataType::Int32,
137        TensorDtype::I64 => DataType::Int64,
138        TensorDtype::U8 => DataType::UInt8,
139        TensorDtype::U32 => DataType::UInt32,
140        TensorDtype::Bool => DataType::Boolean,
141        TensorDtype::F16 => DataType::Float32, // Fallback to F32
142    }
143}
144
145/// Convert an Arrow array to a TensorBlock (zero-copy)
146pub fn arrow_to_tensor_block(array: &dyn Array, shape: TensorShape) -> Result<TensorBlock> {
147    let dtype = arrow_dtype_to_tensor(array.data_type())?;
148
149    // Get the raw buffer data
150    let data = match array.data_type() {
151        DataType::Float32 => {
152            let arr = array
153                .as_any()
154                .downcast_ref::<Float32Array>()
155                .expect("checked: DataType::Float32 matches Float32Array");
156            let buffer = arr.values();
157            // Cast typed slice to &[u8] for Bytes
158            let byte_slice = unsafe {
159                std::slice::from_raw_parts(
160                    buffer.as_ptr() as *const u8,
161                    buffer.len() * std::mem::size_of::<f32>(),
162                )
163            };
164            Bytes::copy_from_slice(byte_slice)
165        }
166        DataType::Float64 => {
167            let arr = array
168                .as_any()
169                .downcast_ref::<Float64Array>()
170                .expect("checked: DataType::Float64 matches Float64Array");
171            let buffer = arr.values();
172            let byte_slice = unsafe {
173                std::slice::from_raw_parts(
174                    buffer.as_ptr() as *const u8,
175                    buffer.len() * std::mem::size_of::<f64>(),
176                )
177            };
178            Bytes::copy_from_slice(byte_slice)
179        }
180        DataType::Int8 => {
181            let arr = array
182                .as_any()
183                .downcast_ref::<Int8Array>()
184                .expect("checked: DataType::Int8 matches Int8Array");
185            let buffer = arr.values();
186            let byte_slice =
187                unsafe { std::slice::from_raw_parts(buffer.as_ptr() as *const u8, buffer.len()) };
188            Bytes::copy_from_slice(byte_slice)
189        }
190        DataType::Int32 => {
191            let arr = array
192                .as_any()
193                .downcast_ref::<Int32Array>()
194                .expect("checked: DataType::Int32 matches Int32Array");
195            let buffer = arr.values();
196            let byte_slice = unsafe {
197                std::slice::from_raw_parts(
198                    buffer.as_ptr() as *const u8,
199                    buffer.len() * std::mem::size_of::<i32>(),
200                )
201            };
202            Bytes::copy_from_slice(byte_slice)
203        }
204        DataType::Int64 => {
205            let arr = array
206                .as_any()
207                .downcast_ref::<Int64Array>()
208                .expect("checked: DataType::Int64 matches Int64Array");
209            let buffer = arr.values();
210            let byte_slice = unsafe {
211                std::slice::from_raw_parts(
212                    buffer.as_ptr() as *const u8,
213                    buffer.len() * std::mem::size_of::<i64>(),
214                )
215            };
216            Bytes::copy_from_slice(byte_slice)
217        }
218        DataType::UInt8 => {
219            let arr = array
220                .as_any()
221                .downcast_ref::<UInt8Array>()
222                .expect("checked: DataType::UInt8 matches UInt8Array");
223            let buffer = arr.values();
224            Bytes::copy_from_slice(buffer.as_ref())
225        }
226        DataType::UInt32 => {
227            let arr = array
228                .as_any()
229                .downcast_ref::<UInt32Array>()
230                .expect("checked: DataType::UInt32 matches UInt32Array");
231            let buffer = arr.values();
232            let byte_slice = unsafe {
233                std::slice::from_raw_parts(
234                    buffer.as_ptr() as *const u8,
235                    buffer.len() * std::mem::size_of::<u32>(),
236                )
237            };
238            Bytes::copy_from_slice(byte_slice)
239        }
240        DataType::Boolean => {
241            let arr = array
242                .as_any()
243                .downcast_ref::<BooleanArray>()
244                .expect("checked: DataType::Boolean matches BooleanArray");
245            let bytes: Vec<u8> = (0..arr.len()).map(|i| arr.value(i) as u8).collect();
246            Bytes::from(bytes)
247        }
248        _ => {
249            return Err(Error::InvalidInput(format!(
250                "Unsupported Arrow dtype: {:?}",
251                array.data_type()
252            )))
253        }
254    };
255
256    TensorBlock::new(data, shape, dtype)
257}
258
259/// Create an Arrow RecordBatch from multiple TensorBlocks
260#[allow(dead_code)]
261pub fn tensors_to_record_batch(
262    tensors: Vec<(&str, &TensorBlock)>,
263) -> Result<arrow_array::RecordBatch> {
264    let mut fields = Vec::new();
265    let mut arrays: Vec<ArrayRef> = Vec::new();
266
267    for (name, tensor) in tensors {
268        fields.push(tensor.to_arrow_field(name));
269        arrays.push(tensor.to_arrow_array()?);
270    }
271
272    let schema = Arc::new(Schema::new(fields));
273    arrow_array::RecordBatch::try_new(schema, arrays)
274        .map_err(|e| Error::InvalidInput(format!("Failed to create RecordBatch: {}", e)))
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    #[test]
282    fn test_tensor_to_arrow_f32() {
283        let data = [1.0f32, 2.0, 3.0, 4.0];
284        let bytes = Bytes::from(
285            data.iter()
286                .flat_map(|&f| f.to_le_bytes())
287                .collect::<Vec<u8>>(),
288        );
289
290        let tensor =
291            TensorBlock::new(bytes, TensorShape::new(vec![2, 2]), TensorDtype::F32).unwrap();
292
293        let arrow_array = tensor.to_arrow_array().unwrap();
294        let f32_array = arrow_array.as_any().downcast_ref::<Float32Array>().unwrap();
295
296        assert_eq!(f32_array.len(), 4);
297        assert_eq!(f32_array.value(0), 1.0);
298        assert_eq!(f32_array.value(1), 2.0);
299        assert_eq!(f32_array.value(2), 3.0);
300        assert_eq!(f32_array.value(3), 4.0);
301    }
302
303    #[test]
304    fn test_arrow_to_tensor_f32() {
305        let arrow_array = Float32Array::from(vec![1.0f32, 2.0, 3.0, 4.0]);
306        let tensor = arrow_to_tensor_block(&arrow_array, TensorShape::new(vec![2, 2])).unwrap();
307
308        assert_eq!(tensor.element_count(), 4);
309        assert_eq!(tensor.metadata().dtype, TensorDtype::F32);
310    }
311
312    #[test]
313    fn test_tensor_to_arrow_i32() {
314        let data = [1i32, 2, 3, 4];
315        let bytes = Bytes::from(
316            data.iter()
317                .flat_map(|&i| i.to_le_bytes())
318                .collect::<Vec<u8>>(),
319        );
320
321        let tensor = TensorBlock::new(bytes, TensorShape::new(vec![4]), TensorDtype::I32).unwrap();
322
323        let arrow_array = tensor.to_arrow_array().unwrap();
324        let i32_array = arrow_array.as_any().downcast_ref::<Int32Array>().unwrap();
325
326        assert_eq!(i32_array.len(), 4);
327        assert_eq!(i32_array.value(0), 1);
328        assert_eq!(i32_array.value(3), 4);
329    }
330
331    #[test]
332    fn test_dtype_conversions() {
333        // TensorDtype to Arrow DataType
334        assert_eq!(tensor_dtype_to_arrow(&TensorDtype::F32), DataType::Float32);
335        assert_eq!(tensor_dtype_to_arrow(&TensorDtype::I64), DataType::Int64);
336        assert_eq!(tensor_dtype_to_arrow(&TensorDtype::Bool), DataType::Boolean);
337
338        // Arrow DataType to TensorDtype
339        assert_eq!(
340            arrow_dtype_to_tensor(&DataType::Float32).unwrap(),
341            TensorDtype::F32
342        );
343        assert_eq!(
344            arrow_dtype_to_tensor(&DataType::Int64).unwrap(),
345            TensorDtype::I64
346        );
347    }
348
349    #[test]
350    fn test_arrow_schema_generation() {
351        let data = Bytes::from(vec![0u8; 16]);
352        let tensor = TensorBlock::new(data, TensorShape::new(vec![4]), TensorDtype::F32).unwrap();
353
354        let schema = tensor.to_arrow_schema("tensor_data");
355        assert_eq!(schema.fields().len(), 1);
356        assert_eq!(schema.field(0).name(), "tensor_data");
357        assert_eq!(schema.field(0).data_type(), &DataType::Float32);
358    }
359
360    #[test]
361    fn test_zero_copy_roundtrip() {
362        // Create Arrow array
363        let original_data = vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0];
364        let arrow_array = Float32Array::from(original_data.clone());
365
366        // Convert to TensorBlock
367        let tensor = arrow_to_tensor_block(&arrow_array, TensorShape::new(vec![2, 3])).unwrap();
368
369        // Convert back to Arrow
370        let arrow_back = tensor.to_arrow_array().unwrap();
371        let f32_back = arrow_back.as_any().downcast_ref::<Float32Array>().unwrap();
372
373        // Verify data integrity
374        assert_eq!(f32_back.len(), original_data.len());
375        for (i, &expected) in original_data.iter().enumerate() {
376            assert_eq!(f32_back.value(i), expected);
377        }
378    }
379
380    #[test]
381    fn test_tensor_to_arrow_field() {
382        let data = Bytes::from(vec![0u8; 64]); // 8 elements * 8 bytes per I64
383        let tensor = TensorBlock::new(data, TensorShape::new(vec![8]), TensorDtype::I64).unwrap();
384
385        let field = tensor.to_arrow_field("my_tensor");
386        assert_eq!(field.name(), "my_tensor");
387        assert_eq!(field.data_type(), &DataType::Int64);
388        assert!(!field.is_nullable());
389    }
390}
ipfrs_core/arrow.rs

ipfrs_core/
arrow.rs
