nordisk/knut

crates/knut-thund/examples/wordcount_stream.rs

//! The canonical streaming example, authored against the Þund IR (the typed
//! Rust builder surface) and lowered to a Spark Declarative Pipelines graph.
//!
//! Run: `cargo run -p knut-thund --example wordcount_stream`
//!
//! This compiles on the default feature set — it exercises authoring, the IR,
//! validation, the SDP lowering, and backend capability checks *without* a
//! live Spark or DataFusion. Wiring an actual run requires `--features spark`
//! (lower onto a real Spark Connect server) or `--features native`.

use knut_thund::backend::{ExecBackend, native::NativeBackend, spark::SparkBackend};
use knut_thund::ir::{
    Dataset, DatasetSchema, Expectation, Flow, FlowKind, OnViolation, OutputType, Pipeline,
    SourceSpec, Trigger, Watermark,
};

fn main() {
    // Author a streaming word-count: read a Kafka topic of lines, tumbling
    // 10s windows, write counts to a streaming table, drop empty words.
    let mut pipeline = Pipeline::new("wordcount")
        .with_dataset(
            Dataset::new("word_counts", OutputType::Table)
                .incremental()
                .with_format("delta")
                .with_schema(
                    DatasetSchema::new()
                        .field("word", "Utf8", false)
                        .field("n", "Int64", false),
                ),
        )
        .with_flow(
            Flow::streaming(
                "count_words",
                "word_counts",
                SourceSpec::Kafka {
                    bootstrap: "localhost:9092".into(),
                    topic: "lines".into(),
                    format: "json".into(),
                },
            )
            .with_watermark(Watermark {
                event_time_column: "ts".into(),
                allowed_lateness_ms: 5_000,
                idle_timeout_ms: 30_000,
            })
            .with_query(
                "SELECT word, count(*) AS n \
                 FROM lines LATERAL VIEW explode(split(line, ' ')) AS word \
                 GROUP BY word, TUMBLE(ts, '10s')",
            )
            .expect(Expectation::new("nonempty_word", "length(word) > 0").on(OnViolation::Drop)),
        );

    // The IR exposes flows directly, so post-hoc tuning is just field
    // mutation — here, set the streaming trigger cadence to 1s micro-batches.
    if let Some(flow) = pipeline.flows.first_mut() {
        if let FlowKind::Streaming { trigger, .. } = &mut flow.kind {
            *trigger = Trigger::ProcessingTime { interval_ms: 1_000 };
        }
    }

    pipeline.validate().expect("pipeline is well-formed");
    println!(
        "pipeline `{}` — streaming: {}",
        pipeline.name,
        pipeline.is_streaming()
    );
    println!("topo order: {:?}", pipeline.topo_order().unwrap());

    // Lower to Spark Declarative Pipelines (graph shape; no live Spark).
    let sdp = pipeline.to_sdp();
    println!(
        "lowered to SDP: {} dataset(s), {} flow(s)",
        sdp.datasets.len(),
        sdp.flows.len()
    );

    // Backend capability checks happen UP FRONT (SanityCheckPlan-style):
    let native = NativeBackend::new().with_state_root("file:///tmp/thund-state");
    match native.check(&pipeline) {
        Ok(()) => println!("native backend: accepts this pipeline"),
        Err(e) => println!("native backend: rejects — {e}"),
    }

    let spark = SparkBackend::new("sc://localhost:15002");
    match spark.check(&pipeline) {
        Ok(()) => println!("spark backend: accepts this pipeline"),
        Err(e) => println!("spark backend: rejects (expected — expectations not in OSS SDP) — {e}"),
    }
}