spiral-sdk 0.1.7

//! Owned Arrow inputs shared by embedded and HTTP table writes.

use arrow_array::{RecordBatch, RecordBatchIterator, RecordBatchReader};
use arrow_schema::SchemaRef;

use crate::Error;
pub use crate::scandal::api::{SortColumn, WriteLayout, WriteProperties};

/// An owned synchronous Arrow reader suitable for bounded streaming uploads.
///
/// Implementations must return from each reader callback; cancelling a write
/// cannot interrupt user code already executing inside a callback.
pub trait IntoArrowStream {
    /// Transfer the schema and batch reader to the SDK.
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error>;

    /// Preserve independent input streams and their physical properties for a table write.
    fn into_write_input(self) -> Result<WriteInput, Error>
    where
        Self: Sized,
    {
        let reader = self.into_arrow_stream()?;
        WriteInput::new(
            reader.schema(),
            vec![reader],
            WriteProperties {
                ordering: Vec::new(),
            },
        )
    }
}

/// Independent Arrow streams with the same schema and guaranteed per-stream ordering.
///
/// Streams upload concurrently, each with bounded buffering. Their fragments publish together;
/// ordering does not imply an order between streams. The storage writer validates ordering.
pub struct WriteInput {
    pub(crate) schema: SchemaRef,
    pub(crate) partitions: Vec<Box<dyn RecordBatchReader + Send>>,
    pub(crate) properties: WriteProperties,
}

impl WriteInput {
    /// Validate the streams' declared schemas and attach their physical properties.
    pub fn new(
        schema: SchemaRef,
        partitions: Vec<Box<dyn RecordBatchReader + Send>>,
        properties: WriteProperties,
    ) -> Result<Self, Error> {
        if partitions
            .iter()
            .any(|partition| partition.schema() != schema)
        {
            return Err(Error::invalid(
                "write partitions must have the declared schema",
            ));
        }
        Ok(Self {
            schema,
            partitions,
            properties,
        })
    }
    /// Shared Arrow schema, also defining an empty write.
    pub fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }
}

impl IntoArrowStream for WriteInput {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        Ok(Box::new(RecordBatchIterator::new(
            self.partitions.into_iter().flatten(),
            self.schema,
        )))
    }
    fn into_write_input(self) -> Result<WriteInput, Error> {
        Ok(self)
    }
}
impl IntoArrowStream for RecordBatch {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        let schema = self.schema();
        Ok(Box::new(RecordBatchIterator::new([Ok(self)], schema)))
    }
}
impl IntoArrowStream for Vec<RecordBatch> {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        let schema = self
            .first()
            .ok_or_else(|| Error::invalid("an empty input needs an explicit schema"))?
            .schema();
        (schema, self).into_arrow_stream()
    }
}
impl IntoArrowStream for (SchemaRef, Vec<RecordBatch>) {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        Ok(Box::new(RecordBatchIterator::new(
            self.1.into_iter().map(Ok),
            self.0,
        )))
    }
}
impl IntoArrowStream for Box<dyn RecordBatchReader + Send> {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        Ok(self)
    }
}
impl<R: RecordBatchReader + Send + 'static> IntoArrowStream for Box<R> {
    fn into_arrow_stream(self) -> Result<Box<dyn RecordBatchReader + Send>, Error> {
        Ok(self)
    }
}