spiral/
arrow_input.rs

1//! Owned Arrow inputs shared by embedded and HTTP table writes.
2
3use arrow_array::{RecordBatch, RecordBatchIterator, RecordBatchReader};
4use arrow_schema::SchemaRef;
5
6use crate::Error;
7pub use crate::scandal::api::{SortColumn, WriteLayout, WriteProperties};
8
9/// An owned synchronous Arrow reader suitable for bounded streaming uploads.
10///
11/// Implementations must return from each reader callback; cancelling a write
12/// cannot interrupt user code already executing inside a callback.
13pub trait IntoArrowStream {
14    /// Transfer the schema and batch reader to the SDK.
15    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error>;
16
17    /// Preserve independent input streams and their physical properties for a table write.
18    fn into_write_input(self) -> Result<WriteInput, Error>
19    where
20        Self: Sized,
21    {
22        let reader = self.into_arrow_stream()?;
23        WriteInput::new(
24            reader.schema(),
25            vec![reader],
26            WriteProperties {
27                ordering: Vec::new(),
28            },
29        )
30    }
31}
32
33/// Independent Arrow streams with the same schema and guaranteed per-stream ordering.
34///
35/// Streams upload concurrently, each with bounded buffering. Their fragments publish together;
36/// ordering does not imply an order between streams. The storage writer validates ordering.
37pub struct WriteInput {
38    pub(crate) schema: SchemaRef,
39    pub(crate) partitions: Vec<Box<dyn RecordBatchReader + Send>>,
40    pub(crate) properties: WriteProperties,
41}
42
43impl WriteInput {
44    /// Validate the streams' declared schemas and attach their physical properties.
45    pub fn new(
46        schema: SchemaRef,
47        partitions: Vec<Box<dyn RecordBatchReader + Send>>,
48        properties: WriteProperties,
49    ) -> Result<Self, Error> {
50        if partitions
51            .iter()
52            .any(|partition| partition.schema() != schema)
53        {
54            return Err(Error::invalid(
55                "write partitions must have the declared schema",
56            ));
57        }
58        Ok(Self {
59            schema,
60            partitions,
61            properties,
62        })
63    }
64    /// Shared Arrow schema, also defining an empty write.
65    pub fn schema(&self) -> SchemaRef {
66        self.schema.clone()
67    }
68}
69
70impl IntoArrowStream for WriteInput {
71    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
72        Ok(Box::new(RecordBatchIterator::new(
73            self.partitions.into_iter().flatten(),
74            self.schema,
75        )))
76    }
77    fn into_write_input(self) -> Result<WriteInput, Error> {
78        Ok(self)
79    }
80}
81impl IntoArrowStream for RecordBatch {
82    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
83        let schema = self.schema();
84        Ok(Box::new(RecordBatchIterator::new([Ok(self)], schema)))
85    }
86}
87impl IntoArrowStream for Vec<RecordBatch> {
88    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
89        let schema = self
90            .first()
91            .ok_or_else(|| Error::invalid("an empty input needs an explicit schema"))?
92            .schema();
93        (schema, self).into_arrow_stream()
94    }
95}
96impl IntoArrowStream for (SchemaRef, Vec<RecordBatch>) {
97    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
98        Ok(Box::new(RecordBatchIterator::new(
99            self.1.into_iter().map(Ok),
100            self.0,
101        )))
102    }
103}
104impl IntoArrowStream for Box<dyn RecordBatchReader + Send> {
105    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
106        Ok(self)
107    }
108}
109impl<R: RecordBatchReader + Send + 'static> IntoArrowStream for Box<R> {
110    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
111        Ok(self)
112    }
113}
spiral/arrow_input.rs

spiral/
arrow_input.rs
