fv_streams_engine/
compute.rs

1//! The worker's compute runtime: a `wasm` / `container` pipeline step selects a Kinetics transform
2//! by `ref`; `fv_plan::KineticsRunner` types the rows, runs it, and returns rows. This module only
3//! decides *which* transforms and backends the worker has: the roots from `FV_TRANSFORMS_DIR`,
4//! and both shipped backends.
5
6use std::path::PathBuf;
7use std::sync::Arc;
8
9use fv_compute::{RegistryError, Runtime};
10use fv_compute_container::ContainerBackend;
11use fv_compute_wasm::WasmBackend;
12pub use fv_plan::KineticsRunner as WorkerCompute;
13
14/// A runner over an ordered list of transform roots (the provided package first, business bundles
15/// after; later roots override by `id@version`), with the wasm and container backends registered.
16pub fn with_roots(roots: &[PathBuf]) -> Result<WorkerCompute, RegistryError> {
17    Ok(WorkerCompute::new(runtime_with_roots(roots)?))
18}
19
20/// The Kinetics runtime itself (registry + backends + cache) over the roots, for the columnar
21/// path that hands a `RecordBatch` to a transform directly.
22pub fn runtime_with_roots(roots: &[PathBuf]) -> Result<Arc<Runtime>, RegistryError> {
23    let mut builder = Runtime::builder();
24    for root in roots {
25        builder = builder.root(root.clone());
26    }
27    let wasm = WasmBackend::new().map_err(|e| RegistryError::Io {
28        path: PathBuf::from("<wasm-engine>"),
29        msg: e.to_string(),
30    })?;
31    Ok(Arc::new(
32        builder.backend(wasm).backend(ContainerBackend::new()).build()?,
33    ))
34}
35
36/// Run a `wasm` / `container` step on a batch: the batch is typed to the transform's declared
37/// input signature (columns by name, cast safely — a value that does not fit lands null, a missing
38/// column is all null), run once, and the transform's output batch returned as is.
39pub fn run_batch(
40    runtime: &Runtime,
41    step: &serde_json::Value,
42    batch: &arrow::array::RecordBatch,
43) -> Result<arrow::array::RecordBatch, String> {
44    let selector = WorkerCompute::selector(step)?;
45    let compute = runtime.load(selector).map_err(|e| e.to_string())?;
46    let manifest = compute.manifest();
47    if manifest.inputs.len() > 1 {
48        return Err(format!(
49            "transform `{selector}` declares {} inputs; a stream stage supplies one",
50            manifest.inputs.len()
51        ));
52    }
53    let typed = match manifest.inputs.first() {
54        Some(signature) => typed_batch(&signature.to_arrow_schema_ref(), batch)?,
55        None => batch.clone(),
56    };
57    compute.run(&[typed]).map_err(|e| e.to_string())
58}
59
60/// `batch` reshaped to `schema`: columns by name, cast safely (unfittable values → null), missing
61/// columns all null.
62pub fn typed_batch(
63    schema: &arrow::datatypes::SchemaRef,
64    batch: &arrow::array::RecordBatch,
65) -> Result<arrow::array::RecordBatch, String> {
66    let n = batch.num_rows();
67    let opts = arrow::compute::CastOptions {
68        safe: true,
69        ..Default::default()
70    };
71    let cols: Vec<arrow::array::ArrayRef> = schema
72        .fields()
73        .iter()
74        .map(|f| match batch.column_by_name(f.name()) {
75            Some(c) if c.data_type() == f.data_type() => Ok(Arc::clone(c)),
76            Some(c) => arrow::compute::cast_with_options(c, f.data_type(), &opts).map_err(|e| e.to_string()),
77            None => Ok(arrow::array::new_null_array(f.data_type(), n)),
78        })
79        .collect::<Result<_, _>>()?;
80    // the signature may declare non-nullable columns; the values decide, so relax the schema.
81    let relaxed = Arc::new(arrow::datatypes::Schema::new(
82        schema
83            .fields()
84            .iter()
85            .map(|f| f.as_ref().clone().with_nullable(true))
86            .collect::<Vec<_>>(),
87    ));
88    arrow::array::RecordBatch::try_new(relaxed, cols).map_err(|e| e.to_string())
89}
90
91/// A runner from `FV_TRANSFORMS_DIR` (colon-separated roots). Unset means an empty registry, so
92/// compute steps fail closed; a malformed manifest in a configured root fails loudly here.
93pub fn from_env() -> Result<WorkerCompute, RegistryError> {
94    let roots: Vec<PathBuf> = std::env::var("FV_TRANSFORMS_DIR")
95        .unwrap_or_default()
96        .split(':')
97        .filter(|s| !s.is_empty())
98        .map(PathBuf::from)
99        .collect();
100    with_roots(&roots)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106    use fv_plan::build::ComputeRunner;
107    use fv_plan::row::Row;
108    use fv_value::Value;
109
110    fn wasm_fixtures() -> PathBuf {
111        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures")
112    }
113
114    #[tokio::test]
115    async fn runs_a_wasm_transform_via_the_registry() {
116        let runner = with_roots(&[wasm_fixtures()]).unwrap();
117        let step = serde_json::json!({ "op": "wasm", "ref": "spikeTotal" });
118        let inputs = vec![(
119            "raw".to_string(),
120            vec![Row(vec![
121                ("id".into(), Value::Num(100.0)),
122                ("amount".into(), Value::Num(7.25)),
123            ])],
124        )];
125        let out = runner.run(&step, &inputs).await.unwrap();
126        assert_eq!(out[0].get("total"), Value::Num(114.5)); // 7.25*2 + 100
127    }
128
129    #[tokio::test]
130    async fn runs_a_container_transform_via_the_registry() {
131        // riskBand: container, json protocol, stdlib python3.
132        let runner = with_roots(&[wasm_fixtures()]).unwrap();
133        let step = serde_json::json!({ "op": "container", "ref": "riskBand" });
134        let inputs = vec![(
135            "staged".to_string(),
136            vec![
137                Row(vec![
138                    ("customerId".into(), Value::Str("C1".into())),
139                    ("creditTerms".into(), Value::Str("NET90".into())),
140                    ("creditLimit".into(), Value::Num(90000.0)),
141                ]),
142                Row(vec![
143                    ("customerId".into(), Value::Str("C2".into())),
144                    ("creditTerms".into(), Value::Str("PREPAID".into())),
145                    ("creditLimit".into(), Value::Num(1000.0)),
146                ]),
147            ],
148        )];
149        let out = runner.run(&step, &inputs).await.unwrap();
150        assert_eq!(out.len(), 2);
151        assert_eq!(out[0].get("customerId"), Value::Str("C1".into()));
152        assert_eq!(out[0].get("riskBand"), Value::Str("HIGH".into())); // NET90 + high limit
153        assert_eq!(out[1].get("riskBand"), Value::Str("LOW".into())); // PREPAID + low limit
154    }
155
156    #[test]
157    fn run_batch_types_the_batch_to_the_signature_and_runs_the_transform() {
158        use arrow::array::{Float64Array, Int64Array, RecordBatch, StringArray};
159        use arrow::datatypes::{DataType, Field, Schema};
160        let rt = runtime_with_roots(&[wasm_fixtures()]).unwrap();
161        let step = serde_json::json!({ "op": "wasm", "ref": "spikeTotal" });
162        // the decoded shape: numbers as Float64, an extra column the signature does not name.
163        let schema = Arc::new(Schema::new(vec![
164            Field::new("id", DataType::Float64, true),
165            Field::new("amount", DataType::Float64, true),
166            Field::new("extra", DataType::Utf8, true),
167        ]));
168        let b = RecordBatch::try_new(
169            schema,
170            vec![
171                Arc::new(Float64Array::from(vec![100.0, 200.0])),
172                Arc::new(Float64Array::from(vec![7.25, 1.0])),
173                Arc::new(StringArray::from(vec!["x", "y"])),
174            ],
175        )
176        .unwrap();
177        let out = run_batch(&rt, &step, &b).unwrap();
178        let total = out.column_by_name("total").unwrap();
179        let totals: Vec<f64> = match total.data_type() {
180            DataType::Float64 => total.as_any().downcast_ref::<Float64Array>().unwrap().values().to_vec(),
181            _ => total
182                .as_any()
183                .downcast_ref::<Int64Array>()
184                .map(|a| a.values().iter().map(|v| *v as f64).collect())
185                .unwrap_or_default(),
186        };
187        assert_eq!(totals, vec![114.5, 202.0]); // amount*2 + id
188    }
189
190    #[tokio::test]
191    async fn unknown_transform_fails_closed() {
192        let runner = with_roots(&[wasm_fixtures()]).unwrap();
193        let step = serde_json::json!({ "op": "wasm", "ref": "doesNotExist" });
194        let err = runner.run(&step, &[("raw".into(), vec![])]).await.unwrap_err();
195        assert!(err.contains("no such transform"), "got: {err}");
196    }
197}
fv_streams_engine/compute.rs

fv_streams_engine/
compute.rs
