spiral/
functions.rs

1//! SPQL functions evaluated by scandal over application-owned Arrow batches.
2
3use std::sync::Arc;
4
5use arrow_array::RecordBatch;
6use arrow_schema::{Schema, SchemaRef};
7use futures::{StreamExt as _, TryStreamExt as _, stream};
8
9use crate::scandal::error::ScandalError;
10use crate::scandal::{Scandal, TableBinding};
11use crate::{BatchStream, Error, FunctionInfo, NamedExpression, Scan};
12
13/// The SPQL functions of one scandal, embedded or remote, and their evaluation.
14///
15/// Scan expressions run inside scans. These calls run the same implementations on the
16/// caller's batches, for eager calls and application-engine UDFs.
17#[derive(Clone)]
18pub struct Functions {
19    scandal: Arc<dyn Scandal>,
20    catalog: Arc<Vec<FunctionInfo>>,
21}
22
23impl Functions {
24    /// Read the catalog of `scandal` once; later calls evaluate with it.
25    pub async fn new(scandal: Arc<dyn Scandal>) -> Result<Self, Error> {
26        let catalog = Arc::new(scandal.list_functions().await?);
27        Ok(Self { scandal, catalog })
28    }
29
30    /// The functions this scandal runs, in identifier order.
31    pub fn list(&self) -> &[FunctionInfo] {
32        &self.catalog
33    }
34
35    /// Bind named expressions to an input schema and return the output schema without data.
36    ///
37    /// Scandal resolves each call's overload and output type exactly as evaluation would.
38    pub async fn bind(
39        &self,
40        outputs: &[NamedExpression],
41        input: &Schema,
42    ) -> Result<SchemaRef, Error> {
43        Ok(self.scandal.bind_projection(input, outputs).await?)
44    }
45
46    /// Evaluate named expressions over a stream of the caller's batches, row for row.
47    ///
48    /// Binding fails before any input is read. One input batch can yield several output
49    /// batches, in order. Dropping the output stops reading the input.
50    pub async fn project_stream(
51        &self,
52        outputs: &[NamedExpression],
53        input: BatchStream,
54    ) -> Result<BatchStream, Error> {
55        self.evaluate_stream(outputs, input, &[]).await
56    }
57
58    /// Project a stream like [`Self::project_stream`], reading blob payloads with the storage
59    /// of the Tables that `scans` read.
60    ///
61    /// A payload under a scan's authorized Table root is read with a fresh binding for that
62    /// Table, and other references fail.
63    pub async fn project_stream_with_scans(
64        &self,
65        outputs: &[NamedExpression],
66        input: BatchStream,
67        scans: &[Scan],
68    ) -> Result<BatchStream, Error> {
69        let mut tables = Vec::with_capacity(scans.len());
70        for scan in scans {
71            tables.extend(scan.table_binding().await?);
72        }
73        self.evaluate_stream(outputs, input, &tables).await
74    }
75
76    /// Evaluate named expressions over one of the caller's batches. The output batches, in
77    /// order, hold one row per input row.
78    pub async fn evaluate(
79        &self,
80        outputs: &[NamedExpression],
81        batch: RecordBatch,
82    ) -> Result<Vec<RecordBatch>, Error> {
83        let input = BatchStream::from_stream(batch.schema(), stream::iter([Ok(batch)]));
84        self.project_stream(outputs, input)
85            .await?
86            .try_collect()
87            .await
88    }
89
90    async fn evaluate_stream(
91        &self,
92        outputs: &[NamedExpression],
93        input: BatchStream,
94        tables: &[TableBinding],
95    ) -> Result<BatchStream, Error> {
96        let schema = input.schema();
97        let input = input
98            .map_err(|error| ScandalError::InvalidArgument {
99                message: error.to_string(),
100            })
101            .boxed();
102        let evaluation = self
103            .scandal
104            .evaluate(outputs, schema, input, tables)
105            .await?;
106        Ok(BatchStream::from_stream(
107            evaluation.schema,
108            evaluation.batches.map_err(Error::from),
109        ))
110    }
111}
spiral/functions.rs

spiral/
functions.rs
