fv_plan/
compute.rs

1//! Compute steps: running a Kinetics transform from a pipeline.
2//!
3//! A `wasm` or `container` step selects a transform from a [`fv_compute::Runtime`] by `ref`
4//! (`id` or `id@version`). [`KineticsRunner`] is the [`ComputeRunner`] that bridges the row-based
5//! build plane to the Arrow-typed transform contract: rows are typed into a `RecordBatch` per
6//! declared input using the manifest's signature, the transform runs, and its output batch comes
7//! back as rows. The runtime's backends decide *how* the transform runs; this module never links
8//! one.
9//!
10//! ```no_run
11//! use std::sync::Arc;
12//! use fv_compute::Runtime;
13//! use fv_plan::{Build, KineticsRunner};
14//! # use fv_compute::{Compute, ComputeError, ImplKind, TransformBackend, TransformManifest};
15//! # struct Noop;
16//! # impl TransformBackend for Noop {
17//! #     fn kind(&self) -> ImplKind { ImplKind::Wasm }
18//! #     fn load(&self, _: &TransformManifest, _: &std::path::Path) -> Result<Box<dyn Compute>, ComputeError> { unimplemented!() }
19//! # }
20//! # let pipeline = serde_json::json!({});
21//! # let datasets: Vec<serde_json::Value> = vec![];
22//!
23//! let runtime = Arc::new(Runtime::builder().root("./transforms").backend(Noop).build()?);
24//! let kinetics = KineticsRunner::new(runtime);
25//!
26//! # tokio::runtime::Runtime::new()?.block_on(async {
27//! let record = Build::new(&pipeline, &datasets).compute(&kinetics).run().await;
28//! # });
29//! # Ok::<(), Box<dyn std::error::Error>>(())
30//! ```
31
32use std::sync::Arc;
33
34use datafusion::arrow::array::{
35    ArrayRef, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray,
36};
37use fv_compute::{FvType, Runtime, SchemaSpec};
38use fv_value::Value;
39use serde_json::Value as J;
40
41use crate::build::ComputeRunner;
42use crate::convert::batch_to_rows;
43use crate::row::Row;
44
45/// Runs `wasm` / `container` steps on a shared [`fv_compute::Runtime`].
46#[derive(Clone)]
47pub struct KineticsRunner {
48    runtime: Arc<Runtime>,
49}
50
51impl KineticsRunner {
52    pub fn new(runtime: Arc<Runtime>) -> Self {
53        Self { runtime }
54    }
55
56    /// The runtime this runner dispatches to.
57    pub fn runtime(&self) -> &Arc<Runtime> {
58        &self.runtime
59    }
60
61    /// The transform a compute step selects: its `ref` (or legacy `transform`) field.
62    pub fn selector(step: &J) -> Result<&str, String> {
63        step["ref"]
64            .as_str()
65            .or_else(|| step["transform"].as_str())
66            .ok_or_else(|| "compute step missing `ref` (the transform id or id@version)".to_string())
67    }
68
69    /// Resolve and load the step's transform without running it. Call this when a pipeline or
70    /// stream is configured, so a missing or broken transform fails there rather than per batch.
71    pub fn ensure_loadable(&self, step: &J) -> Result<(), String> {
72        self.runtime
73            .ensure_loadable(Self::selector(step)?)
74            .map_err(|e| e.to_string())
75    }
76
77    /// Run a compute step synchronously. The transform backends are synchronous (a component
78    /// call, or a blocking child process), so this needs no async runtime; the async
79    /// [`ComputeRunner`] impl is a thin wrapper.
80    pub fn run_sync(&self, step: &J, inputs: &[(String, Vec<Row>)]) -> Result<Vec<Row>, String> {
81        let selector = Self::selector(step)?;
82        let compute = self.runtime.load(selector).map_err(|e| e.to_string())?;
83        let manifest = compute.manifest();
84        if !manifest.inputs.is_empty() && manifest.inputs.len() != inputs.len() {
85            return Err(format!(
86                "transform `{selector}` declares {} input(s), the pipeline supplied {}",
87                manifest.inputs.len(),
88                inputs.len()
89            ));
90        }
91        let batches: Vec<RecordBatch> = manifest
92            .inputs
93            .iter()
94            .zip(inputs.iter())
95            .map(|(signature, (_, rows))| rows_to_typed_batch(signature, rows))
96            .collect::<Result<_, _>>()?;
97        let out = compute.run(&batches).map_err(|e| e.to_string())?;
98        Ok(batch_to_rows(&out))
99    }
100}
101
102#[async_trait::async_trait]
103impl ComputeRunner for KineticsRunner {
104    async fn run(&self, step: &J, inputs: &[(String, Vec<Row>)]) -> Result<Vec<Row>, String> {
105        self.run_sync(step, inputs)
106    }
107}
108
109/// Type rows into a `RecordBatch` by a declared signature (a manifest input), rather than by
110/// inference: each column gets the Arrow type the signature names, and a value that does not fit
111/// lands null.
112pub fn rows_to_typed_batch(signature: &SchemaSpec, rows: &[Row]) -> Result<RecordBatch, String> {
113    let arrays: Vec<ArrayRef> = signature
114        .columns
115        .iter()
116        .map(|c| typed_column(c.dtype, &c.name, rows))
117        .collect::<Result<_, _>>()?;
118    RecordBatch::try_new(signature.to_arrow_schema_ref(), arrays).map_err(|e| e.to_string())
119}
120
121fn number(v: &Value) -> Option<f64> {
122    match v {
123        Value::Num(n) => Some(*n),
124        _ => None,
125    }
126}
127
128fn text(v: &Value) -> Option<String> {
129    match v {
130        Value::Null => None,
131        Value::Str(s) => Some(s.clone()),
132        Value::Bool(b) => Some(if *b { "true" } else { "false" }.into()),
133        Value::Num(n) => Some(if n.fract() == 0.0 && n.is_finite() {
134            format!("{}", *n as i64)
135        } else {
136            format!("{n}")
137        }),
138        other => Some(format!("{other:?}")),
139    }
140}
141
142fn typed_column(dtype: FvType, name: &str, rows: &[Row]) -> Result<ArrayRef, String> {
143    let cell = |r: &Row| r.get(name);
144    Ok(match dtype {
145        FvType::Int64 => Arc::new(
146            rows.iter()
147                .map(|r| number(&cell(r)).map(|n| n as i64))
148                .collect::<Int64Array>(),
149        ),
150        FvType::Int32 => Arc::new(
151            rows.iter()
152                .map(|r| number(&cell(r)).map(|n| n as i32))
153                .collect::<Int32Array>(),
154        ),
155        FvType::Float64 => Arc::new(rows.iter().map(|r| number(&cell(r))).collect::<Float64Array>()),
156        FvType::Float32 => Arc::new(
157            rows.iter()
158                .map(|r| number(&cell(r)).map(|n| n as f32))
159                .collect::<Float32Array>(),
160        ),
161        FvType::Bool => Arc::new(
162            rows.iter()
163                .map(|r| match cell(r) {
164                    Value::Bool(b) => Some(b),
165                    _ => None,
166                })
167                .collect::<BooleanArray>(),
168        ),
169        FvType::Utf8 => Arc::new(rows.iter().map(|r| text(&cell(r))).collect::<StringArray>()),
170        other => {
171            return Err(format!(
172                "compute step: column `{name}` has type {other:?}, which the row bridge does not carry yet"
173            ))
174        }
175    })
176}
fv_plan/compute.rs

fv_plan/
compute.rs
