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