Skip to main content

df_interchange/
to_arrow.rs

1use paste::paste;
2use std::mem::transmute;
3
4use crate::{error::InterchangeError, ArrowArray, ArrowSchema, Interchange};
5
6macro_rules! ffi_to_arrow {
7    ($to_ver:literal) => {
8        paste! {
9            impl Interchange {
10                #[doc = "Move Arrow data interchange format to Arrow version `" $to_ver "`."]
11                pub fn [<to_arrow_ $to_ver>](mut self) -> Result<Vec<[<arrow_crate_ $to_ver>]::record_batch::RecordBatch>, InterchangeError> {
12
13                    if !self.chunks_aligned {
14                        return Err(InterchangeError::ChunksNotAligned);
15                    }
16
17                    // Get number of columns
18                    let num_cols = self.ffi.len();
19
20                    // Get number of chunks
21                    let num_chunks = self.ffi[0].1.len();
22
23                    // Get batch vector ready
24                    let mut batches = Vec::with_capacity(num_chunks);
25
26                    // For columns in the ffi
27                    for _ in 0..num_chunks {
28
29                        // Create arrays for the batch
30                        let mut arrays = Vec::with_capacity(num_cols);
31                        let mut fields = Vec::with_capacity(num_cols);
32
33                        for col_num in 0..num_cols {
34
35                            // Get the chunk
36                            let chunk = self.ffi[col_num].1.pop().unwrap();
37
38                            // Convert ffi array from this crate's version of ArrowArray to arrow-rs
39                            let ffi_array = unsafe { transmute::<ArrowArray,  [<arrow_crate_ $to_ver>]::ffi::FFI_ArrowArray>(chunk.0) };
40
41                            // Convert ffi field from this crate's version of ArrowField to arrow-rs
42                            let ffi_schema = unsafe { transmute::<ArrowSchema,  [<arrow_crate_ $to_ver>]::ffi::FFI_ArrowSchema>(chunk.1) };
43
44                            // Import into arrow-rs
45                            let from_ffi = unsafe { [<arrow_crate_ $to_ver>]::ffi::from_ffi(ffi_array, &ffi_schema) }?;
46
47                            // Make an array out of it
48                            let array_ref = [<arrow_crate_ $to_ver>]::array::make_array(from_ffi);
49
50                            // Get the field for the batch schema
51                            let field = std::convert::TryInto::<[<arrow_crate_ $to_ver>]::datatypes::Field>::try_into(&ffi_schema)?;
52
53                            arrays.push(array_ref);
54                            fields.push(std::sync::Arc::new(field));
55                        }
56
57                        // Create batch record from array and schema
58                        let schema = [<arrow_crate_ $to_ver>]::datatypes::Schema::new(fields);
59                        let record_batch = [<arrow_crate_ $to_ver>]::record_batch::RecordBatch::try_new(std::sync::Arc::new(schema), arrays)?;
60
61                        batches.push(record_batch)
62
63                    }
64
65                    // Reverse the batches
66                    let batches = batches.into_iter().rev().collect();
67
68                    Ok(batches)
69                }
70            }
71        }
72    };
73}
74
75#[cfg(feature = "arrow_54")]
76ffi_to_arrow!("54");
77
78#[cfg(feature = "arrow_55")]
79ffi_to_arrow!("55");
80
81#[cfg(feature = "arrow_56")]
82ffi_to_arrow!("56");
83
84#[cfg(feature = "arrow_57")]
85ffi_to_arrow!("57");