fv-streams-engine 0.6.0

//! The worker's compute runtime: a `wasm` / `container` pipeline step selects a Kinetics transform
//! by `ref`; `fv_plan::KineticsRunner` types the rows, runs it, and returns rows. This module only
//! decides *which* transforms and backends the worker has: the roots from `FV_TRANSFORMS_DIR`,
//! and both shipped backends.

use std::path::PathBuf;
use std::sync::Arc;

use fv_compute::{RegistryError, Runtime};
use fv_compute_container::ContainerBackend;
use fv_compute_wasm::WasmBackend;
pub use fv_plan::KineticsRunner as WorkerCompute;

/// A runner over an ordered list of transform roots (the provided package first, business bundles
/// after; later roots override by `id@version`), with the wasm and container backends registered.
pub fn with_roots(roots: &[PathBuf]) -> Result<WorkerCompute, RegistryError> {
    Ok(WorkerCompute::new(runtime_with_roots(roots)?))
}

/// The Kinetics runtime itself (registry + backends + cache) over the roots, for the columnar
/// path that hands a `RecordBatch` to a transform directly.
pub fn runtime_with_roots(roots: &[PathBuf]) -> Result<Arc<Runtime>, RegistryError> {
    let mut builder = Runtime::builder();
    for root in roots {
        builder = builder.root(root.clone());
    }
    let wasm = WasmBackend::new().map_err(|e| RegistryError::Io {
        path: PathBuf::from("<wasm-engine>"),
        msg: e.to_string(),
    })?;
    Ok(Arc::new(
        builder.backend(wasm).backend(ContainerBackend::new()).build()?,
    ))
}

/// Run a `wasm` / `container` step on a batch: the batch is typed to the transform's declared
/// input signature (columns by name, cast safely — a value that does not fit lands null, a missing
/// column is all null), run once, and the transform's output batch returned as is.
pub fn run_batch(
    runtime: &Runtime,
    step: &serde_json::Value,
    batch: &arrow::array::RecordBatch,
) -> Result<arrow::array::RecordBatch, String> {
    let selector = WorkerCompute::selector(step)?;
    let compute = runtime.load(selector).map_err(|e| e.to_string())?;
    let manifest = compute.manifest();
    if manifest.inputs.len() > 1 {
        return Err(format!(
            "transform `{selector}` declares {} inputs; a stream stage supplies one",
            manifest.inputs.len()
        ));
    }
    let typed = match manifest.inputs.first() {
        Some(signature) => typed_batch(&signature.to_arrow_schema_ref(), batch)?,
        None => batch.clone(),
    };
    compute.run(&[typed]).map_err(|e| e.to_string())
}

/// `batch` reshaped to `schema`: columns by name, cast safely (unfittable values → null), missing
/// columns all null.
pub fn typed_batch(
    schema: &arrow::datatypes::SchemaRef,
    batch: &arrow::array::RecordBatch,
) -> Result<arrow::array::RecordBatch, String> {
    let n = batch.num_rows();
    let opts = arrow::compute::CastOptions {
        safe: true,
        ..Default::default()
    };
    let cols: Vec<arrow::array::ArrayRef> = schema
        .fields()
        .iter()
        .map(|f| match batch.column_by_name(f.name()) {
            Some(c) if c.data_type() == f.data_type() => Ok(Arc::clone(c)),
            Some(c) => arrow::compute::cast_with_options(c, f.data_type(), &opts).map_err(|e| e.to_string()),
            None => Ok(arrow::array::new_null_array(f.data_type(), n)),
        })
        .collect::<Result<_, _>>()?;
    // the signature may declare non-nullable columns; the values decide, so relax the schema.
    let relaxed = Arc::new(arrow::datatypes::Schema::new(
        schema
            .fields()
            .iter()
            .map(|f| f.as_ref().clone().with_nullable(true))
            .collect::<Vec<_>>(),
    ));
    arrow::array::RecordBatch::try_new(relaxed, cols).map_err(|e| e.to_string())
}

/// A runner from `FV_TRANSFORMS_DIR` (colon-separated roots). Unset means an empty registry, so
/// compute steps fail closed; a malformed manifest in a configured root fails loudly here.
pub fn from_env() -> Result<WorkerCompute, RegistryError> {
    let roots: Vec<PathBuf> = std::env::var("FV_TRANSFORMS_DIR")
        .unwrap_or_default()
        .split(':')
        .filter(|s| !s.is_empty())
        .map(PathBuf::from)
        .collect();
    with_roots(&roots)
}

#[cfg(test)]
mod tests {
    use super::*;
    use fv_plan::build::ComputeRunner;
    use fv_plan::row::Row;
    use fv_value::Value;

    fn wasm_fixtures() -> PathBuf {
        PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures")
    }

    #[tokio::test]
    async fn runs_a_wasm_transform_via_the_registry() {
        let runner = with_roots(&[wasm_fixtures()]).unwrap();
        let step = serde_json::json!({ "op": "wasm", "ref": "spikeTotal" });
        let inputs = vec![(
            "raw".to_string(),
            vec![Row(vec![
                ("id".into(), Value::Num(100.0)),
                ("amount".into(), Value::Num(7.25)),
            ])],
        )];
        let out = runner.run(&step, &inputs).await.unwrap();
        assert_eq!(out[0].get("total"), Value::Num(114.5)); // 7.25*2 + 100
    }

    #[tokio::test]
    async fn runs_a_container_transform_via_the_registry() {
        // riskBand: container, json protocol, stdlib python3.
        let runner = with_roots(&[wasm_fixtures()]).unwrap();
        let step = serde_json::json!({ "op": "container", "ref": "riskBand" });
        let inputs = vec![(
            "staged".to_string(),
            vec![
                Row(vec![
                    ("customerId".into(), Value::Str("C1".into())),
                    ("creditTerms".into(), Value::Str("NET90".into())),
                    ("creditLimit".into(), Value::Num(90000.0)),
                ]),
                Row(vec![
                    ("customerId".into(), Value::Str("C2".into())),
                    ("creditTerms".into(), Value::Str("PREPAID".into())),
                    ("creditLimit".into(), Value::Num(1000.0)),
                ]),
            ],
        )];
        let out = runner.run(&step, &inputs).await.unwrap();
        assert_eq!(out.len(), 2);
        assert_eq!(out[0].get("customerId"), Value::Str("C1".into()));
        assert_eq!(out[0].get("riskBand"), Value::Str("HIGH".into())); // NET90 + high limit
        assert_eq!(out[1].get("riskBand"), Value::Str("LOW".into())); // PREPAID + low limit
    }

    #[test]
    fn run_batch_types_the_batch_to_the_signature_and_runs_the_transform() {
        use arrow::array::{Float64Array, Int64Array, RecordBatch, StringArray};
        use arrow::datatypes::{DataType, Field, Schema};
        let rt = runtime_with_roots(&[wasm_fixtures()]).unwrap();
        let step = serde_json::json!({ "op": "wasm", "ref": "spikeTotal" });
        // the decoded shape: numbers as Float64, an extra column the signature does not name.
        let schema = Arc::new(Schema::new(vec![
            Field::new("id", DataType::Float64, true),
            Field::new("amount", DataType::Float64, true),
            Field::new("extra", DataType::Utf8, true),
        ]));
        let b = RecordBatch::try_new(
            schema,
            vec![
                Arc::new(Float64Array::from(vec![100.0, 200.0])),
                Arc::new(Float64Array::from(vec![7.25, 1.0])),
                Arc::new(StringArray::from(vec!["x", "y"])),
            ],
        )
        .unwrap();
        let out = run_batch(&rt, &step, &b).unwrap();
        let total = out.column_by_name("total").unwrap();
        let totals: Vec<f64> = match total.data_type() {
            DataType::Float64 => total.as_any().downcast_ref::<Float64Array>().unwrap().values().to_vec(),
            _ => total
                .as_any()
                .downcast_ref::<Int64Array>()
                .map(|a| a.values().iter().map(|v| *v as f64).collect())
                .unwrap_or_default(),
        };
        assert_eq!(totals, vec![114.5, 202.0]); // amount*2 + id
    }

    #[tokio::test]
    async fn unknown_transform_fails_closed() {
        let runner = with_roots(&[wasm_fixtures()]).unwrap();
        let step = serde_json::json!({ "op": "wasm", "ref": "doesNotExist" });
        let err = runner.run(&step, &[("raw".into(), vec![])]).await.unwrap_err();
        assert!(err.contains("no such transform"), "got: {err}");
    }
}