Datum

Appearance

Cookbook ​

Practical, copy-paste recipes. Most code blocks are imported directly from compiled, asserting Rust and Python tests under crates/*/tests/; all projection snippets use that mechanism. The remaining inline persistence recipes mirror the shipped integration-test API and keep their shared serializer helper visible on the page.

Recipes ​

Rate-limit a stream ​

Problem: cap outbound calls to at most N requests per unit of time.

Use throttle(max_elements, per, max_burst, ThrottleMode) to pace the stream. ThrottleMode::Shaping introduces delay; ThrottleMode::Enforcing fails excess elements with StreamError::Failed.

rust
use datum::{Sink, Source, ThrottleMode};
use std::time::{Duration, Instant};

// Problem: cap outbound API calls to at most 1 request every 20 ms.
// throttle(max_elements, per, max_burst, mode) paces the stream.
// ThrottleMode::Shaping delays excess elements; Enforcing fails them instead.
let t0 = Instant::now();
let items: Vec<u64> = Source::from_iter(1_u64..=5)
    .throttle(1, Duration::from_millis(20), 1, ThrottleMode::Shaping)
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();
let elapsed = t0.elapsed();

Watch: max_burst controls the initial token bucket — set it equal to max_elements for a strict no-burst policy.

Keep-latest under backpressure ​

Problem: a slow consumer receives sensor bursts; merge them rather than buffering forever.

conflate(f) coalesces buffered elements when the consumer is slower than the producer. |_acc, latest| latest discards older values and retains the most-recent arrival.

rust
use datum::{Sink, Source};

// Problem: a slow consumer receives sensor bursts; keep only the most-recent value.
// conflate(f) coalesces buffered elements when the consumer is slower than the
// producer.  "Keep-latest" discards the accumulated value, retaining the newest.
let items: Vec<u64> = Source::from_iter(1_u64..=5)
    .conflate(|_accumulated, latest| latest)
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();

Watch: in a synchronous fused chain the source outruns the consumer before the first pull, so all values collapse into one. Use conflate(|acc, x| acc + x) when you need to preserve a summary (e.g. a running total) across the burst.

Recover from upstream failure with a fallback source ​

Problem: a primary data source may fail transiently; continue with a static default.

recover_with(f) intercepts the first error and replaces it with a Source<T> continuation. Elements from the fallback are emitted normally and the stream completes cleanly.

rust
use datum::{Sink, Source, StreamError};

// Problem: a primary data source may fail; fall back to a static default dataset
// so downstream processing can continue.
// recover_with(f) replaces the first error with a Source<T> continuation.
let items: Vec<u64> = Source::from_iter(1_u64..=2)
    .concat(Source::failed(StreamError::Failed(
        "upstream unavailable".into(),
    )))
    .recover_with(|_err| Some(Source::from_iter([99_u64, 100])))
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();

Watch: recover_with_retries(n, f) re-attempts the fallback source up to n times — useful when the fallback itself may fail. recover(f) is the simpler variant that emits at most one fallback element.

Frame a byte stream ​

Problem: a TCP connection delivers log lines in arbitrary-sized chunks.

Framing::delimiter(delimiter, max_frame_len, allow_truncation) reassembles frames across chunk boundaries and strips the delimiter. allow_truncation = true emits a partial final frame when the stream ends without a trailing delimiter.

rust
use datum::{Framing, Sink, Source};

// Problem: a TCP connection delivers log lines in arbitrary-sized byte chunks.
// Framing::delimiter(delimiter, max_frame_len, allow_truncation) reassembles
// frames across chunk boundaries and strips the delimiter.
// allow_truncation=true emits a partial final frame with no trailing delimiter.
let chunks: Vec<Vec<u8>> = vec![
    b"event:login\nevent:".to_vec(), // chunk boundary splits a frame
    b"logout\n".to_vec(),
    b"event:timeout".to_vec(), // no trailing newline
];

let frames: Vec<Vec<u8>> = Source::from_iter(chunks)
    .via(Framing::delimiter(b"\n".to_vec(), 256, true))
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();

Watch: max_frame_len protects against unbounded buffering on a malformed stream — exceeding it fails with StreamError::Failed. Framing::json(max_len) is available for concatenated JSON objects.

Resilient source with exponential-backoff restart ​

Problem: a remote source fails intermittently; reconnect automatically with backoff.

RestartSource::on_failures_with_backoff(settings, factory) rematerializes the factory when the inner source fails. RestartSettings::with_max_restarts(n, within) bounds the total restart count inside a rolling window.

rust
use datum::{RestartSettings, RestartSource, Sink, Source, StreamError};
use std::sync::{
    Arc,
    atomic::{AtomicUsize, Ordering},
};
use std::time::Duration;

let attempts = Arc::new(AtomicUsize::new(0));
let attempts_c = Arc::clone(&attempts);

let settings = RestartSettings::new(
    Duration::from_millis(5),   // min backoff
    Duration::from_millis(100), // max backoff
    0.0,                        // no jitter — deterministic in tests
)
.with_max_restarts(3, Duration::from_secs(10));

// on_failures_with_backoff rematerializes the factory when the inner source fails;
// a clean completion passes through unchanged.  Useful for reconnecting to external
// services that experience transient errors.
let items: Vec<u64> = RestartSource::on_failures_with_backoff(settings, move || {
    let n = attempts_c.fetch_add(1, Ordering::SeqCst);
    if n == 0 {
        // First attempt: simulate a transient upstream error.
        Source::failed(StreamError::Failed("transient error".into()))
    } else {
        // Second attempt succeeds.
        Source::from_iter(1_u64..=3)
    }
})
.take(3)
.run_with(Sink::collect())
.unwrap()
.wait()
.unwrap();

Watch: on_failures_with_backoff only restarts on failure — a clean completion flows through unchanged. Use with_backoff if you want to restart on both failure and completion (e.g. for infinite pagination).

Fan out and recombine with GraphDSL ​

Problem: apply two independent transformations to each element, then merge results.

The GraphDSL builder wires typed ports together. Broadcast fans one inlet to N outlets; Merge collects N inlets into one outlet.

rust
use datum::{Broadcast, GraphDsl, GraphFlowShape, Merge};

// Problem: apply two independent processing paths to each element, then merge
// the results.  GraphDSL wires arbitrary graph topologies via typed ports.
// Here Broadcast fans every element to two outlets; Merge collects both.
let graph = GraphDsl::try_create(|builder| {
    let bcast = builder.add(Broadcast::<u64>::new(2));
    let merge = builder.add(Merge::<u64>::new(2));
    builder.connect(bcast.outlet(0)?, merge.inlet(0)?)?;
    builder.connect(bcast.outlet(1)?, merge.inlet(1)?)?;
    Ok(GraphFlowShape::new(bcast.inlet(), merge.outlet()))
})
.unwrap();

// Each input element passes through both outlets → every value appears twice.
let mut items: Vec<u64> = graph.run_with_input(1_u64..=3).unwrap();
items.sort_unstable();

Watch: Merge is non-deterministic when inlets deliver concurrently — sort before asserting equality. For a round-robin fan-in use Balance; for an ordered fan-in use MergeSequence.

Group-by aggregation ​

Problem: compute per-category totals from a mixed event stream.

group_by(max_keys, key_fn, allow_recreation) partitions elements into per-key sub-sources. flat_map_merge(breadth, f) runs up to breadth sub-source folds concurrently.

rust
use datum::{Sink, Source};

// Problem: compute per-category totals from a mixed event stream.
// group_by(max_keys, key_fn, allow_recreation) partitions elements by key.
// Each sub-source is folded independently; flat_map_merge runs the folds
// concurrently up to max_keys at a time.
let events: Vec<(u32, u64)> = vec![(1, 10), (2, 20), (1, 30), (2, 40), (1, 50)];

let mut totals: Vec<u64> = Source::from_iter(events)
    .group_by(2, |e| e.0, false)
    .flat_map_merge(2, |sub| sub.fold(0_u64, |acc, e| acc + e.1))
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();
totals.sort_unstable();

Watch: max_keys caps the number of simultaneously open sub-sources; elements with an unknown key after the cap is reached will fail the stream.

Bounded async parallelism ​

Problem: enrich records with async lookups without overwhelming the backend.

map_async(parallelism, f) runs up to parallelism futures concurrently and preserves input order in the output. Use map_async_unordered when order does not matter and you want lower latency.

rust
use datum::{Sink, Source};

// Problem: enrich each record with an async lookup (e.g. a DB or HTTP call),
// running at most 4 concurrent requests to avoid saturating the backend.
// map_async(parallelism, f) caps concurrent futures and preserves input order.
let enriched: Vec<String> = Source::from_iter(1_u64..=6)
    .map_async(4, |id| async move { Ok(format!("user-{id}")) })
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();

Watch: the future closure must return StreamResult<Next> (i.e. Result<Next, StreamError>). A returned Err fails the whole stream; wrap recoverable errors explicitly if you need per-element fallback logic.

Ask an actor per element ​

Problem: delegate per-element computation to a typed Ractor actor.

ActorFlow::ask(actor_ref, parallelism, timeout, make_msg) routes each stream element as a request/reply pair. Up to parallelism requests may be in-flight simultaneously; output order is preserved.

Happy path:

rust
use datum::actor::{Actor, ActorProcessingErr, ActorRef};
use datum::{ActorFlow, ReplyPort, Sink, Source};
use std::time::Duration;

enum SquareMsg {
    Square {
        value: u64,
        reply_to: ReplyPort<u64>,
    },
}

struct SquareActor;

impl Actor for SquareActor {
    type Msg = SquareMsg;
    type State = ();
    type Arguments = ();

    async fn pre_start(
        &self,
        _: ActorRef<Self::Msg>,
        _: Self::Arguments,
    ) -> Result<Self::State, ActorProcessingErr> {
        Ok(())
    }

    async fn handle(
        &self,
        _: ActorRef<Self::Msg>,
        message: Self::Msg,
        _: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        let SquareMsg::Square { value, reply_to } = message;
        let _ = reply_to.send(value * value);
        Ok(())
    }
}

// Problem: delegate per-element computation to a typed actor.
// ActorFlow::ask(actor_ref, parallelism, timeout, make_msg) routes each element
// as a request/reply pair and preserves output order.
let rt = tokio::runtime::Builder::new_multi_thread()
    .enable_all()
    .build()
    .unwrap();
let (actor_ref, _handle) = rt.block_on(async {
    Actor::spawn(None, SquareActor, ())
        .await
        .expect("actor spawns")
});

let results: Vec<u64> = Source::from_iter(1_u64..=5)
    .via(ActorFlow::ask(
        actor_ref.clone(),
        2, // at most 2 in-flight requests
        Duration::from_secs(1),
        |v, reply_to| SquareMsg::Square { value: v, reply_to },
    ))
    .run_with(Sink::collect())
    .unwrap()
    .wait()
    .unwrap();

actor_ref.stop(None);

Timeout / failure path — when the actor never replies within the deadline the stream fails with StreamError::ActorAskTimeout:

rust
use datum::actor::{Actor, ActorProcessingErr, ActorRef};
use datum::{ActorFlow, ReplyPort, Sink, Source, StreamError};
use std::time::Duration;

// Problem: the backend actor is overloaded and never replies within the deadline.
// ActorFlow::ask propagates the missed deadline as StreamError::ActorAskTimeout,
// so callers can distinguish a timeout from an actor crash.
enum HangMsg {
    Hang { _reply_to: ReplyPort<u64> },
}

struct HangActor;

impl Actor for HangActor {
    type Msg = HangMsg;
    type State = Vec<ReplyPort<u64>>;
    type Arguments = ();

    async fn pre_start(
        &self,
        _: ActorRef<Self::Msg>,
        _: Self::Arguments,
    ) -> Result<Self::State, ActorProcessingErr> {
        Ok(Vec::new())
    }

    async fn handle(
        &self,
        _: ActorRef<Self::Msg>,
        message: Self::Msg,
        state: &mut Self::State,
    ) -> Result<(), ActorProcessingErr> {
        let HangMsg::Hang { _reply_to } = message;
        state.push(_reply_to); // hold port open — never reply
        Ok(())
    }
}

let rt = tokio::runtime::Builder::new_multi_thread()
    .enable_all()
    .build()
    .unwrap();
let (actor_ref, _handle) = rt.block_on(async {
    Actor::spawn(None, HangActor, ())
        .await
        .expect("actor spawns")
});

let timeout = Duration::from_millis(50);
let result = Source::single(1_u64)
    .via(ActorFlow::ask(
        actor_ref.clone(),
        1,
        timeout,
        |_v, reply_to| HangMsg::Hang {
            _reply_to: reply_to,
        },
    ))
    .run_with(Sink::collect())
    .unwrap()
    .wait();

actor_ref.stop(None);

Watch: actor-ask throughput carries ~2× CPU vs wall-clock on the ordered_sum benchmark path — the ask protocol does real work. See Actors Interop for the full performance picture.

Cancel a running stream ​

Problem: shut down a long-running stream from an external handle.

KillSwitches::single() materializes a UniqueKillSwitch alongside the stream. switch.shutdown() completes downstream and cancels upstream at the next pull. switch.abort(err) fails the stream instead.

rust
use datum::{
    Keep, KillSwitches, Materializer,
    testkit::{TestSink, TestSource},
};

// Problem: shut down a long-running stream from an external handle without
// stopping the whole runtime.
// KillSwitches::single() materializes a UniqueKillSwitch alongside the stream.
// switch.shutdown() completes downstream and cancels upstream on the next pull.
let materializer = Materializer::new();
let ((source, switch), sink) = TestSource::probe::<u64>()
    .via_mat(KillSwitches::single(), Keep::both)
    .to_mat(TestSink::probe(), Keep::both)
    .run_with_materializer(&materializer)
    .expect("materializes");

// Deliver one element normally before triggering shutdown.
sink.request(1);
assert_eq!(source.expect_request(), 1);
source.send_next(42);
sink.assert_next(42);

// Shutdown: the next pull observes completion; upstream receives cancellation.
switch.shutdown();
sink.request(1);
sink.expect_complete();
source.expect_cancellation();

Watch: KillSwitches::shared(name) produces a SharedKillSwitch whose switch.flow::<T>() can be wired into any number of independent streams — one shutdown() call terminates all of them.

Dynamic fan-in from many producers ​

Problem: aggregate input from an unknown number of concurrent producers that attach at different times.

MergeHub::source_with_draining::<T>(buffer) materializes a Sink<T, NotUsed> that any producer can clone and attach to independently. control.drain_and_complete() signals that no new producers will attach.

rust
use datum::{Keep, MergeHub, Sink, Source};

// Problem: aggregate input from an unknown number of concurrent producers
// that attach at different times.
// MergeHub::source_with_draining materializes a Sink<T, NotUsed> that any
// producer can clone and attach to independently.
let ((hub_sink, control), collected) = MergeHub::source_with_draining::<u32>(4)
    .to_mat(Sink::collect(), Keep::both)
    .run()
    .unwrap();

// Each producer runs its own source into a clone of the hub sink.
hub_sink
    .clone()
    .run_with(Source::from_iter([1_u32, 2, 3]))
    .unwrap();
hub_sink
    .clone()
    .run_with(Source::from_iter([10_u32, 20, 30]))
    .unwrap();

// Signal that no new producers will attach; hub completes when all finish.
control.drain_and_complete();

let mut items: Vec<u32> = collected.wait().unwrap();
items.sort_unstable();

Watch: per_producer_buffer_size is per-producer, not global. Keep it small to bound memory; producers that fill their buffer are backpressured.

Drive and assert a pipeline with TestKit probes ​

Problem: verify a processing pipeline's element-by-element behavior in isolation.

TestSource::probe::<T>() and TestSink::probe::<T>() give manual control over every upstream element and downstream credit signal. The protocol is explicit: credit must be issued before elements can arrive.

rust
use datum::{
    Flow, Keep, Materializer,
    testkit::{TestSink, TestSource},
};

// Problem: verify a processing pipeline's element-by-element behavior in
// isolation before wiring it into production.
// TestSource and TestSink probes give manual control over upstream elements
// and downstream credit, making the full backpressure protocol visible.
let materializer = Materializer::new();
let (source, sink) = TestSource::probe::<u64>()
    .via(Flow::identity().map(|x| format!("id-{x}")))
    .to_mat(TestSink::probe(), Keep::both)
    .run_with_materializer(&materializer)
    .expect("materializes");

// One credit → one pull propagates upstream → one element travels end-to-end.
sink.request(1);
assert_eq!(source.expect_request(), 1);
source.send_next(42);
sink.assert_next("id-42".to_string());

sink.request(1);
assert_eq!(source.expect_request(), 1);
source.send_next(99);
sink.assert_next("id-99".to_string());

// Complete the publisher; subscriber observes completion on the next pull.
sink.request(1);
assert_eq!(source.expect_request(), 1);
source.send_complete();
sink.expect_complete();

Watch: terminal signals (expect_complete, expect_error) require outstanding credit — always call sink.request(n) before asserting a completion or error, or use drain_until_complete() to consume all elements in one call.

Consume Kafka with the native connector ​

Problem: consume a Kafka topic with Datum-managed at-least-once commits, using the native Kafka connector.

The snippet keeps KafkaConsumerBackend::Native explicit, uses CommitPolicy::Manual, and commits only after your side effect or checkpoint succeeds. It is broker-gated: set MQ_BOOTSTRAP_SERVERS and DATUM_DOCS_ORDERS_TOPIC to run it.

rust
use datum::{Sink, StreamError};
use datum_mq::{
    CommitPolicy, ConsumerRecord, KafkaConsumerBackend, KafkaConsumerSettings, KafkaOffset,
    KafkaSource, Subscription,
};

// Broker-gated: set these when running the recipe against a local Kafka topic
// that already contains at least three order records.
let Some(topic) = std::env::var("DATUM_DOCS_ORDERS_TOPIC").ok() else {
    return;
};
let bootstrap =
    std::env::var("MQ_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "127.0.0.1:9092".to_owned());
let group = format!("datum-docs-native-{}", std::process::id());

let settings = KafkaConsumerSettings::new(bootstrap, group)
    .with_consumer_backend(KafkaConsumerBackend::Native)
    .with_commit_policy(CommitPolicy::Manual)
    .with("auto.offset.reset", "earliest")
    .with("partition.assignment.strategy", "cooperative-sticky")
    .with_backpressure(512, 1_024)
    .with_poll_batch_size(64);

let processed = KafkaSource::committable(settings, Subscription::topics([topic]))
    .as_source()
    .take(3)
    .run_with(Sink::fold_result(
        0_u64,
        |count, (record, offset): (ConsumerRecord, KafkaOffset)| {
            // Do durable, idempotent work first. Use topic/partition/offset
            // from the record or offset as your de-duplication key.
            let _dedupe_key = (&record.topic, record.partition, record.offset);

            // Commit only after the side effect/checkpoint succeeds.
            offset.commit().map_err(StreamError::from)?;
            Ok(count + 1)
        },
    ))
    .expect("native Kafka consumer materializes")
    .wait()
    .expect("native Kafka consumer completes");

println!("processed {processed} orders");
assert_eq!(processed, 3);

Expected output:

text
processed 3 orders

Watch: Kafka is native-only in Datum 0.11: there is no rdkafka feature or dependency, and no C/CMake build. The native producer and consumer support plaintext, ring-backed TLS, PLAIN/SCRAM SASL, and none/gzip/Snappy/LZ4/Zstd record batches. Transactions, consume-transform-produce EOS, GSSAPI, and OAUTHBEARER are not claimed.

Run a native Kafka consumer group ​

Problem: run multiple native consumers in one group, use cooperative-sticky assignment, stop cleanly, and avoid committing offsets after an assignment is lost.

The native backend uses classic consumer groups. Cooperative-sticky is the default, but setting partition.assignment.strategy=cooperative-sticky makes the rollout obvious. The snippet handles AssignmentLost as a terminal commit error and uses KafkaControl::drain_and_shutdown for deploy shutdown. Set MQ_BOOTSTRAP_SERVERS and DATUM_DOCS_GROUP_TOPIC to run it.

rust
use std::time::Duration;

use datum::{Keep, Sink, StreamCompletion, StreamError};
use datum_mq::{
    KafkaConsumerBackend, KafkaConsumerSettings, KafkaControl, KafkaPayloadBatch, KafkaSource,
    MqError, Subscription,
};

// Broker-gated: the topic should exist and receive records while this
// process is running.
let Some(topic) = std::env::var("DATUM_DOCS_GROUP_TOPIC").ok() else {
    return;
};
let bootstrap =
    std::env::var("MQ_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "127.0.0.1:9092".to_owned());

let settings = KafkaConsumerSettings::new(bootstrap, "orders-workers")
    .with_consumer_backend(KafkaConsumerBackend::Native)
    .with("auto.offset.reset", "earliest")
    // The native backend accepts one classic assignor. Cooperative-sticky is
    // the default; set it explicitly when you want rollout parity.
    .with("partition.assignment.strategy", "cooperative-sticky")
    .with_backpressure(256, 512)
    .with_commit_batch_size(128)
    .with_commit_interval(Duration::from_millis(100));

let (control, completion): (KafkaControl, StreamCompletion<datum::NotUsed>) =
    KafkaSource::committable_payload_batches(settings, Subscription::topics([topic]))
        .to_mat(
            Sink::foreach_result(|batch: KafkaPayloadBatch| {
                for record in batch.records() {
                    let _payload = batch.payload(record);
                    // Process each payload here.
                }

                match batch.commit() {
                    Ok(()) => Ok(()),
                    Err(error @ MqError::AssignmentLost { .. }) => {
                        // The group no longer owns at least one partition in
                        // this batch. Do not retry this commit; fail the
                        // stream and let the group replay from the last
                        // committed offset.
                        Err(StreamError::from(error))
                    }
                    Err(error) => Err(StreamError::from(error)),
                }
            }),
            Keep::both,
        )
        .run()
        .expect("native Kafka group consumer materializes");

// On SIGTERM or deploy shutdown, stop taking new records and wait for
// already emitted batches to commit.
control
    .drain_and_shutdown(Duration::from_secs(30))
    .expect("consumer drains outstanding commits");
completion.wait().expect("consumer completes after drain");

let metrics = control.metrics().snapshot();
println!(
    "drained; committed={} rebalances={} lost={}",
    metrics.committed_offsets, metrics.rebalances, metrics.lost_partitions
);

Expected output:

text
drained; committed=<offset-count> rebalances=<count> lost=0

Watch: AssignmentLost means the group no longer owns that partition. Do not retry that commit; fail or stop the stream and let Kafka replay from the last committed offset on the new owner.

Run SQL over Kafka with native ingest ​

Problem: expose a Kafka JSON topic as a SQL table and execute a streaming SQL query through datum-sql, while consuming Kafka through the native connector.

register_mq_topic builds a committable SQL table over datum-mq payload batches. Kafka ingest uses Datum's native connector; the snippet keeps with_consumer_backend(KafkaConsumerBackend::Native) explicit alongside the consumer settings. Set MQ_BOOTSTRAP_SERVERS and DATUM_DOCS_SQL_TOPIC to run the snippet against JSON rows shaped like {"id":201,"name":"q0-a"}.

rust
use std::sync::{Arc, Mutex};

use arrow::array::Int64Array;
use datum_mq::{KafkaConsumerBackend, KafkaConsumerSettings};
use datum_sql::{DatumSqlContext, JsonRowFormat};

// Broker-gated: point this at a Kafka topic containing JSON rows shaped like
// {"id":201,"name":"q0-a"}.
let Some(topic) = std::env::var("DATUM_DOCS_SQL_TOPIC").ok() else {
    return;
};
let bootstrap =
    std::env::var("MQ_BOOTSTRAP_SERVERS").unwrap_or_else(|_| "127.0.0.1:9092".to_owned());

let schema = test_schema();
let format = JsonRowFormat::new(Arc::clone(&schema)).with_schema_revision(18);
let settings = KafkaConsumerSettings::new(
    bootstrap,
    format!("datum-docs-sql-native-{}", unique_suffix()),
)
.with_consumer_backend(KafkaConsumerBackend::Native)
.with("auto.offset.reset", "earliest")
.with_poll_batch_size(128);

let context = DatumSqlContext::new();
context
    .register_mq_topic("orders", settings, topic, format)
    .expect("Kafka topic registers as a SQL table");

let ids = Arc::new(Mutex::new(Vec::<i64>::new()));
let ids_for_sink = Arc::clone(&ids);
context
    .register_append_sink("selected_orders", move |batch| {
        let id_column = batch
            .column(0)
            .as_any()
            .downcast_ref::<Int64Array>()
            .expect("id is Int64");
        ids_for_sink
            .lock()
            .expect("ids lock")
            .extend((0..id_column.len()).map(|row| id_column.value(row)));
        Ok(())
    })
    .expect("append sink registers");

let handle = context
    .execute_streaming(
        "INSERT INTO selected_orders \
         SELECT id FROM orders WHERE id >= 200",
    )
    .await
    .expect("streaming SQL query materializes");

let deadline = Instant::now() + Duration::from_secs(30);
while ids.lock().expect("ids lock").len() < 3 && Instant::now() < deadline {
    std::thread::sleep(Duration::from_millis(50));
}
handle.cancel();

let mut ids = ids.lock().expect("ids lock").clone();
ids.sort_unstable();
println!("selected order ids: {ids:?}");

Expected output:

text
selected order ids: [201, 202, 203]

Watch: execute_streaming keeps running until the source completes, fails, or is cancelled. Keep the returned handle and call cancel() or wire it into your service shutdown path.

Respawn remembered entities after rebalance ​

Problem: keep entity ids alive across shard movement so a rebalanced shard starts its remembered entities on the new owner without waiting for a triggering message.

Enable RememberEntitiesConfig with a store, register the type through register_remembered_entity_type, and call flush_remember_entities() when you need a shutdown/restart-visible barrier for queued start/stop records. The v0.10.3 behavior guarantee is that allocation replication is retried during shard movement, so remembered entity re-spawn follows the final owner after rebalance.

rust
use std::{
    collections::BTreeMap,
    net::SocketAddr,
    sync::{Arc, Mutex},
};

use datum_agent::{
    ClusterAgent, ClusterAgentConfig, ClusterAgentHandle, NodeSessionConfig,
    dcp::{DcpJobFactories, DcpServerConfig, DcpTcpServerConfig},
};
use datum_cluster::{ClusterConfig, MemberState};
use datum_cluster_sharding::{
    EntityContext, InMemoryStore, RememberEntitiesConfig, ReplyPort, Sharding, ShardingConfig,
    ShardingHandle, ShardingResult,
};

#[derive(Serialize, Deserialize)]
enum CartMsg {
    Touch(ReplyPort<String>),
}

fn agent_config(node_id: &str, seed_nodes: Vec<SocketAddr>) -> ClusterAgentConfig {
    ClusterAgentConfig {
        cluster: ClusterConfig {
            node_id: node_id.to_owned(),
            seed_nodes,
            bind_addr: "127.0.0.1:0".parse().expect("bind addr"),
            advertise_addr: "127.0.0.1:0".parse().expect("advertise addr"),
            gossip_interval: Duration::from_millis(60),
            probe_timeout: Duration::from_millis(15),
            downing_timeout: Duration::from_millis(150),
            ..ClusterConfig::default()
        },
        dcp: DcpServerConfig {
            node_id: node_id.to_owned(),
            tcp: Some(DcpTcpServerConfig {
                addr: "127.0.0.1:0".parse().expect("dcp addr"),
            }),
            ..DcpServerConfig::default()
        },
        sessions: NodeSessionConfig {
            request_timeout: Duration::from_secs(2),
            command_buffer: 256,
            ..NodeSessionConfig::default()
        },
        ..ClusterAgentConfig::default()
    }
}

async fn wait_all_up(agents: &[&ClusterAgentHandle], expected: usize) {
    let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
    while tokio::time::Instant::now() < deadline {
        if agents.iter().all(|agent| {
            agent
                .cluster()
                .current_state()
                .members
                .values()
                .filter(|member| member.state == MemberState::Up && !member.unreachable)
                .count()
                == expected
        }) {
            return;
        }
        tokio::time::sleep(Duration::from_millis(20)).await;
    }
    panic!("cluster did not converge to {expected} Up members");
}

async fn register_cart(
    sharding: &ShardingHandle,
    starts: Arc<Mutex<BTreeMap<String, Vec<String>>>>,
) {
    let node_id = sharding.node_id().to_owned();
    sharding
        .register_remembered_entity_type("cart", move |context: EntityContext| {
            starts
                .lock()
                .expect("starts lock")
                .entry(context.entity_id.clone())
                .or_default()
                .push(node_id.clone());
            let reply_node = node_id.clone();
            move |_context: &EntityContext, message: CartMsg| -> ShardingResult<()> {
                let CartMsg::Touch(reply) = message;
                let _ = reply.send(reply_node.clone());
                Ok(())
            }
        })
        .await
        .expect("remembered entity type registers");
}

let starts = Arc::new(Mutex::new(BTreeMap::<String, Vec<String>>::new()));
let remember_store = Arc::new(InMemoryStore::new());
let sharding_config = ShardingConfig {
    num_shards: 16,
    rebalance_per_round: 16,
    coordinator_tick: Duration::from_millis(25),
    remember_entities: RememberEntitiesConfig::with_store(remember_store),
    ..ShardingConfig::default()
};

let first_agent =
    ClusterAgent::start(agent_config("node-a", Vec::new()), DcpJobFactories::new())
        .await
        .expect("first node starts");
let first = Sharding::init(&first_agent, sharding_config.clone()).expect("first sharding");
register_cart(&first, Arc::clone(&starts)).await;

// Start enough entity ids to cover many shards, then flush the write-behind
// remember store before inducing a rebalance.
let mut entity_shards = BTreeMap::new();
for index in 0..64 {
    let entity_id = format!("cart-{index}");
    let entity_ref = first.entity_ref::<CartMsg>("cart", entity_id.clone());
    entity_shards.insert(entity_id, entity_ref.shard_id().to_owned());
    let _owner = entity_ref
        .ask(Duration::from_secs(5), CartMsg::Touch)
        .await
        .expect("initial cart touch");
}
first
    .flush_remember_entities()
    .await
    .expect("remembered starts flush");

let initial_table = first
    .allocation_table("cart")
    .await
    .expect("initial allocation table");
let initial_owner_by_shard = initial_table
    .entries
    .iter()
    .map(|entry| (entry.shard_id.clone(), entry.node_id.clone()))
    .collect::<BTreeMap<_, _>>();

// A joining node triggers graceful shard rebalance. Remembered ids for any
// moved shard are re-spawned on the new owner without a new user message.
let second_agent = ClusterAgent::start(
    agent_config("node-b", vec![first_agent.cluster().advertise_addr()]),
    DcpJobFactories::new(),
)
.await
.expect("second node starts");
wait_all_up(&[&first_agent, &second_agent], 2).await;
let second = Sharding::init(&second_agent, sharding_config).expect("second sharding");
register_cart(&second, Arc::clone(&starts)).await;

let moved_entity = {
    let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
    loop {
        let table = first
            .allocation_table("cart")
            .await
            .expect("current allocation table");
        let final_owner_by_shard = table
            .entries
            .iter()
            .map(|entry| (entry.shard_id.clone(), entry.node_id.clone()))
            .collect::<BTreeMap<_, _>>();
        if let Some(entity_id) = entity_shards.iter().find_map(|(entity_id, shard_id)| {
            let before = initial_owner_by_shard.get(shard_id)?;
            let after = final_owner_by_shard.get(shard_id)?;
            (before != after).then(|| entity_id.clone())
        }) {
            break entity_id;
        }
        if tokio::time::Instant::now() >= deadline {
            panic!("no remembered cart moved during rebalance");
        }
        tokio::time::sleep(Duration::from_millis(25)).await;
    }
};

let deadline = tokio::time::Instant::now() + Duration::from_secs(10);
while tokio::time::Instant::now() < deadline {
    if starts
        .lock()
        .expect("starts lock")
        .get(&moved_entity)
        .is_some_and(|nodes| nodes.len() >= 2)
    {
        break;
    }
    tokio::time::sleep(Duration::from_millis(25)).await;
}

let started_on = starts
    .lock()
    .expect("starts lock")
    .get(&moved_entity)
    .cloned()
    .expect("moved entity was started");
println!("{moved_entity} started on {started_on:?}");
assert!(started_on.len() >= 2);

first.shutdown().await;
second.shutdown().await;
first_agent.shutdown().await.expect("first shuts down");
second_agent.shutdown().await.expect("second shuts down");

Expected output:

text
cart-17 started on ["node-a", "node-b"]

Watch: remember-entities remembers ids, not in-memory state. A moved or restarted entity is constructed from its factory again; use your own persistence for durable counters, carts, sessions, or aggregates.

Build and run a Python pipeline ​

Problem: assemble and run a Datum stream from Python using plain Python callables.

map(lambda ...) and filter(lambda ...) are convenience operators on integer streams. Datum batches elements into single-column Arrow RecordBatch values and executes the generated Arrow UDF once per batch; the row loop runs inside Python. The graph remains a blueprint until run(runtime).

python
import threading
import time

import pytest

import datum
import datum._native as native


def test_fold_sum_over_range_matches_rust_result():
    with datum.Runtime() as runtime:
        graph = (
            datum.Source.range(0, 1_000)
            .map_multiply(3)
            .filter_greater_than(1_000)
            .to_mat(datum.Sink.fold())
        )

        assert graph.run(runtime).wait() == native._rust_pipeline_sum(1_000)


def test_named_fold_terminals_preserve_sum_and_product_defaults():
    with datum.Runtime() as runtime:
        source = datum.Source.from_iter([2, 3, 4])
        assert source.to_mat(datum.Sink.fold()).run(runtime).wait() == 9
        assert source.to_mat(datum.Sink.fold_sum(10)).run(runtime).wait() == 19
        assert source.to_mat(datum.Sink.fold_product()).run(runtime).wait() == 24
        assert source.to_mat(datum.Sink.fold_product(2)).run(runtime).wait() == 48


@pytest.mark.parametrize(
    ("op", "replacement"),
    [
        ("sum", "fold_sum"),
        ("add", "fold_sum"),
        ("+", "fold_sum"),
        ("product", "fold_product"),
        ("mul", "fold_product"),
        ("multiply", "fold_product"),
        ("*", "fold_product"),
    ],
)
def test_removed_fold_op_names_exact_replacement(op, replacement):
    with pytest.raises(datum.BuildError) as error:
        datum.Sink.fold(7, op=op)

    assert str(error.value) == (
        f"fold() no longer accepts op=; use {replacement}(7)"
    )


def test_removed_fold_op_unknown_value_names_both_typed_replacements():
    with pytest.raises(datum.BuildError) as error:
        datum.Sink.fold(7, "custom")

    assert str(error.value) == (
        "fold() no longer accepts op=; use fold_sum(initial) or "
        "fold_product(initial)"
    )


def test_blueprint_can_be_materialized_twice():
    runtime = datum.Runtime()
    graph = (
        datum.Source.from_iter([1, 2, 3, 4])
        .map_add(10)
        .drop(1)
        .take(2)
        .to_mat(datum.Sink.collect())
    )

    try:
        assert graph.run(runtime).wait() == [12, 13]
        assert graph.run(runtime).wait() == [12, 13]
    finally:
        runtime.shutdown()


def test_stream_error_preserves_rust_text():
    with datum.Runtime() as runtime:
        graph = datum.Source.single(1).map_add(2**63 - 1).to_mat(datum.Sink.collect())

        with pytest.raises(datum.StreamError, match="integer overflow in map_add"):
            graph.run(runtime).wait()


def test_named_map_and_filter_kernels_build_integer_pipelines():
    with datum.Runtime() as runtime:
        graph = (
            datum.Source.from_iter([1, 2, 3])
            .map_add(10)
            .filter_greater_than(11)
            .to_mat(datum.Sink.collect())
        )

        assert graph.run(runtime).wait() == [12, 13]

        flow = datum.Flow.map_multiply(2).via(datum.Flow.filter_less_or_equal(4))
        flow_graph = datum.Source.from_iter([1, 2, 3]).via(flow).to_mat(datum.Sink.collect())

        assert flow_graph.run(runtime).wait() == [2, 4]


@pytest.mark.parametrize(
    ("op", "replacement"),
    [
        ("add", "map_add"),
        ("+", "map_add"),
        ("sub", "map_subtract"),
        ("subtract", "map_subtract"),
        ("-", "map_subtract"),
        ("mul", "map_multiply"),
        ("multiply", "map_multiply"),
        ("*", "map_multiply"),
    ],
)
def test_removed_source_map_op_names_exact_replacement(op, replacement):
    with pytest.raises(datum.BuildError) as error:
        datum.Source.single(1).map_op(op, 3)

    assert str(error.value) == (
        f"map_op() was removed; use {replacement}(3), map(lambda) for arbitrary "
        "Python, or col()-expressions on batch streams"
    )


@pytest.mark.parametrize(
    ("op", "replacement"),
    [
        ("eq", "filter_equal"),
        ("equal", "filter_equal"),
        ("==", "filter_equal"),
        ("ne", "filter_not_equal"),
        ("not_equal", "filter_not_equal"),
        ("!=", "filter_not_equal"),
        ("lt", "filter_less_than"),
        ("less_than", "filter_less_than"),
        ("<", "filter_less_than"),
        ("le", "filter_less_or_equal"),
        ("less_or_equal", "filter_less_or_equal"),
        ("<=", "filter_less_or_equal"),
        ("gt", "filter_greater_than"),
        ("greater_than", "filter_greater_than"),
        (">", "filter_greater_than"),
        ("ge", "filter_greater_or_equal"),
        ("greater_or_equal", "filter_greater_or_equal"),
        (">=", "filter_greater_or_equal"),
    ],
)
def test_removed_source_filter_op_names_exact_replacement(op, replacement):
    with pytest.raises(datum.BuildError) as error:
        datum.Source.single(1).filter_op(op, 3)

    assert str(error.value) == (
        f"filter_op() was removed; use {replacement}(3), filter(lambda) for arbitrary "
        "Python, or col()-expressions on batch streams"
    )


def test_removed_flow_string_ops_and_unknown_ops_raise_generic_messages():
    with pytest.raises(datum.BuildError) as map_error:
        datum.Flow.map_op("python_udf", 1)
    assert str(map_error.value) == (
        "map_op() was removed; use map_add(), map_subtract(), or map_multiply(), "
        "map(lambda) for arbitrary Python, or col()-expressions on batch streams"
    )

    with pytest.raises(datum.BuildError) as filter_error:
        datum.Flow.filter_op("python_udf", 1)
    assert str(filter_error.value) == (
        "filter_op() was removed; use filter_equal(), filter_not_equal(), "
        "filter_less_than(), filter_less_or_equal(), filter_greater_than(), or "
        "filter_greater_or_equal(), filter(lambda) for arbitrary Python, or "
        "col()-expressions on batch streams"
    )


def test_map_filter_and_flat_map_accept_python_callables():
    with datum.Runtime() as runtime:
        graph = (
            datum.Source.from_iter([1, 2, 3, 4, 5])
            .map(lambda item: item + 1, batch_size=2)
            .filter(lambda item: item > 3, batch_size=2)
            .flat_map(lambda item: [item, item * 10], batch_size=2)
            .to_mat(datum.Sink.collect())
        )

        assert graph.run(runtime).wait() == [4, 40, 5, 50, 6, 60]

        flow = datum.Flow.map(lambda item: item * 2, batch_size=3).via(
            datum.Flow.filter(lambda item: item >= 4, batch_size=3)
        )
        flow = flow.via(datum.Flow.flat_map(lambda item: (item, item + 1), batch_size=3))
        assert (
            datum.Source.from_iter([1, 2, 3])
            .via(flow)
            .to_mat(datum.Sink.collect())
            .run(runtime)
            .wait()
        ) == [4, 5, 6, 7]

    source = datum.Source.single(1)
    assert not hasattr(source, "map_arithmetic")
    assert not hasattr(source, "filter_compare")
    assert not hasattr(datum.Flow, "map_arithmetic")
    assert not hasattr(datum.Flow, "filter_compare")


def test_callable_type_contract_failures_preserve_traceback_text():
    with datum.Runtime() as runtime:
        with pytest.raises(datum.StreamError) as map_error:
            (
                datum.Source.single(1)
                .map(lambda item: "not an int")
                .to_mat(datum.Sink.collect())
                .run(runtime)
                .wait()
            )
        assert "Traceback" in str(map_error.value)
        assert "map callable result must be int" in str(map_error.value)

        with pytest.raises(datum.StreamError) as filter_error:
            (
                datum.Source.single(1)
                .filter(lambda item: 1)
                .to_mat(datum.Sink.collect())
                .run(runtime)
                .wait()
            )
        assert "filter callable must return bool" in str(filter_error.value)

        with pytest.raises(datum.StreamError) as flat_map_error:
            (
                datum.Source.single(1)
                .flat_map(lambda item: 1)
                .to_mat(datum.Sink.collect())
                .run(runtime)
                .wait()
            )
        assert "flat_map callable must return an iterable of int" in str(
            flat_map_error.value
        )


def test_build_errors_use_datum_exception_hierarchy():
    assert issubclass(datum.BuildError, datum.DatumError)
    assert issubclass(datum.StreamError, datum.DatumError)


def test_generic_runtime_aliases_are_subscriptable():
    assert datum.Source[int] is datum.Source
    assert datum.Flow[int, int] is datum.Flow
    assert datum.Sink[int, list[int]] is datum.Sink
    assert datum.RunnableGraph[int] is datum.RunnableGraph
    assert datum.StreamCompletion[int] is datum.StreamCompletion
    assert datum.Inlet[int] is datum.Inlet
    assert datum.Outlet[int] is datum.Outlet


def test_stream_completion_wait_is_one_shot():
    with datum.Runtime() as runtime:
        completion = datum.Source.single(1).to_mat(datum.Sink.collect()).run(runtime)

        assert completion.wait() == [1]
        with pytest.raises(datum.BuildError) as error:
            completion.wait()

    assert "already consumed" in str(error.value)
    assert "cached result reuse is intentionally unsupported" in str(error.value)


def test_runtime_context_manager_shuts_down():
    runtime = datum.Runtime()
    with runtime as entered:
        assert entered is runtime
        assert not runtime.is_shutdown()

    assert runtime.is_shutdown()


def test_wait_releases_gil_while_rust_runs():
    progress = 0
    running = True

    def spin():
        nonlocal progress
        while running:
            progress += 1
            time.sleep(0)

    worker = threading.Thread(target=spin)
    worker.start()
    with datum.Runtime() as runtime:
        graph = (
            datum.Source.range(0, 20_000_000)
            .map_multiply(3)
            .filter_greater_than(20_000_000)
            .to_mat(datum.Sink.fold())
        )
        completion = graph.run(runtime)
        before = progress
        assert completion.wait() == native._rust_pipeline_sum(20_000_000)
        after = progress

    running = False
    worker.join(timeout=5)
    assert after > before


def test_for_each_count_and_flow_via():
    with datum.Runtime() as runtime:
        flow = datum.Flow.map_add(1).via(datum.Flow.filter_greater_than(2))
        graph = datum.Source.from_iter([1, 2, 3]).via(flow).to_mat(datum.Sink.for_each_count())

        assert graph.run(runtime).wait() == 2


def test_cookbook_python_pipeline_recipe():
    # #region cookbook-python-pipeline
    import datum

    with datum.Runtime() as runtime:
        graph = (
            datum.Source.from_iter([1, 2, 3, 4])
            .map(lambda item: item + 10)
            .filter(lambda item: item > 12)
            .to_mat(datum.Sink.collect())
        )
        result = graph.run(runtime).wait()

    print(result)
    assert result == [13, 14]
    # #endregion cookbook-python-pipeline

Expected output:

text
[13, 14]

Watch: the callable tier is ergonomic, not the fastest path. Use named Rust kernels such as map_add / filter_greater_than for hot integer arithmetic and comparisons, or map_batches with PyArrow kernels for vectorized Python work. Use Sink.fold_sum() or Sink.fold_product() for named arithmetic terminals, and PartitionStrategy.MODULO when a GraphDSL partition strategy is explicit.

Run a vectorized Arrow UDF from Python ​

Problem: transform Arrow batches with a vectorized Python function and surface Python failures as clean Datum stream errors.

Source.from_arrow(...).map_batches(udf, output_schema=...) passes whole pyarrow.RecordBatch values across the boundary. Declare the output schema up front; Datum validates the first returned batch and preserves Python traceback text when a UDF raises.

python
import threading
import time

import pytest

import datum
import datum._native as native

pa = pytest.importorskip("pyarrow")
pc = pytest.importorskip("pyarrow.compute")


def batch(values=(1, 2, 3, 4)):
    ids = pa.array(values, type=pa.int64())
    offsets = pa.array([value * 10 for value in values], type=pa.int64())
    return pa.record_batch([ids, offsets], names=["id", "offset"])


def one_column_schema(name="id"):
    return pa.schema([(name, pa.int64())])


def empty_batch(schema=None):
    schema = schema or one_column_schema()
    arrays = [pa.array([], type=field.type) for field in schema]
    return pa.record_batch(arrays, schema=schema)


def buffer_address(record_batch, column=0, buffer=1):
    return record_batch.column(column).buffers()[buffer].address


def test_source_from_table_reader_and_iterator_count_rows():
    first = batch((1, 2))
    second = batch((3, 4, 5))

    with datum.Runtime() as runtime:
        table = pa.Table.from_batches([first, second])
        assert datum.Source.from_arrow(table).count_rows(runtime) == 5

        reader = pa.RecordBatchReader.from_batches(first.schema, [first, second])
        assert datum.Source.from_arrow(reader).count_rows(runtime) == 5

        assert datum.Source.from_arrow(iter([first, second])).count_rows(runtime) == 5


def test_source_from_arrow_empty_input_requires_declared_schema():
    schema = one_column_schema()

    with pytest.raises(datum.BuildError, match="schema is required for an empty Arrow batch iterator"):
        datum.Source.from_arrow([])

    with pytest.raises(datum.BuildError, match="schema is required for empty Arrow input"):
        datum.Source.from_arrow(empty_batch(schema))

    with pytest.raises(datum.BuildError, match="schema is required for empty Arrow input"):
        datum.Source.from_arrow(pa.Table.from_batches([empty_batch(schema)]))

    source = datum.Source.from_arrow([], schema=schema)
    assert source.schema() == schema
    with datum.Runtime() as runtime:
        assert source.count_rows(runtime) == 0


def test_source_from_arrow_declared_schema_validates_at_build_time():
    source_batch = batch()
    declared = one_column_schema()

    with pytest.raises(datum.BuildError, match="input Arrow object schema mismatch"):
        datum.Source.from_arrow(source_batch, schema=declared)

    mismatched = pa.record_batch(
        [pa.array([1, 2], type=pa.int64())],
        names=["other"],
    )
    with pytest.raises(datum.BuildError, match="input Arrow batch schema mismatch"):
        datum.Source.from_arrow([batch((1, 2)), mismatched])


def test_identity_udf_round_trips_zero_copy():
    source_batch = batch()

    with datum.Runtime() as runtime:
        result = (
            datum.Source.from_arrow(source_batch)
            .map_batches(lambda incoming: incoming, output_schema=source_batch.schema)
            .run_collect(runtime)
        )

    result_batch = result.to_batches()[0]
    assert result_batch.to_pydict() == source_batch.to_pydict()
    assert buffer_address(result_batch) == buffer_address(source_batch)

    input_buffers, output_buffers = native._debug_arrow_udf_return_buffer_addresses(
        source_batch,
        lambda incoming: incoming,
        source_batch.schema,
        identity_fast_path=False,
    )
    assert output_buffers == input_buffers


def test_projection_udf_declares_and_validates_output_schema():
    source_batch = batch()
    output_schema = one_column_schema()

    def project_id(incoming):
        return pa.record_batch([incoming.column(0)], schema=output_schema)

    with datum.Runtime() as runtime:
        result = (
            datum.Source.from_arrow(source_batch)
            .map_batches(project_id, output_schema=output_schema)
            .run_collect(runtime)
        )

    assert result.schema == output_schema
    assert result.to_pydict() == {"id": [1, 2, 3, 4]}
    assert buffer_address(result.to_batches()[0]) == buffer_address(source_batch)

    input_buffers, output_buffers = native._debug_arrow_udf_return_buffer_addresses(
        source_batch,
        project_id,
        output_schema,
    )
    assert output_buffers == input_buffers[:1]


def test_batch_schema_propagates_through_sources_and_flows():
    source_batch = batch()
    projected_schema = one_column_schema()
    doubled_schema = pa.schema([("doubled", pa.int64())])

    source = datum.Source.from_arrow(source_batch)
    assert source.schema() == source_batch.schema

    flow = datum.Flow.map_batches(
        lambda incoming: pa.record_batch([incoming.column(0)], schema=projected_schema),
        output_schema=projected_schema,
    )
    assert flow.schema() == projected_schema

    chained_flow = flow.then_map_batches(
        lambda incoming: pa.record_batch(
            [pc.multiply(incoming.column(0), pa.scalar(2, type=pa.int64()))],
            schema=doubled_schema,
        ),
        output_schema=doubled_schema,
    )
    assert chained_flow.schema() == doubled_schema
    assert source.via(chained_flow).schema() == doubled_schema


def test_map_batches_requires_declared_output_schema_at_build_time():
    source_batch = batch()
    source = datum.Source.from_arrow(source_batch)

    with pytest.raises(datum.BuildError, match="requires declared output_schema"):
        source.map_batches(lambda incoming: incoming)

    with pytest.raises(datum.BuildError, match="requires declared output_schema"):
        datum.Flow.map_batches(lambda incoming: incoming)

    flow = datum.Flow.map_batches(
        lambda incoming: incoming,
        output_schema=source_batch.schema,
    )
    with pytest.raises(datum.BuildError, match="requires declared output_schema"):
        flow.then_map_batches(lambda incoming: incoming)


def test_compute_kernel_udf_uses_arrow_batches():
    source_batch = batch()
    output_schema = one_column_schema("total")

    def add_columns(incoming):
        total = pc.add(incoming.column(0), incoming.column(1))
        return pa.record_batch([total], schema=output_schema)

    with datum.Runtime() as runtime:
        result = (
            datum.Source.from_arrow(source_batch)
            .map_batches(add_columns, output_schema=output_schema)
            .run_collect(runtime)
        )

    assert result.to_pydict() == {"total": [11, 22, 33, 44]}


def test_udf_return_schema_mismatch_fails_stream_at_first_batch():
    source_batch = batch()
    output_schema = one_column_schema()

    stream = datum.Source.from_arrow(source_batch).map_batches(
        lambda incoming: incoming,
        output_schema=output_schema,
    )

    with datum.Runtime() as runtime:
        with pytest.raises(datum.StreamError, match="schema mismatch"):
            stream.count_rows(runtime)


def test_batch_source_via_rejects_non_batch_flow_at_build_call():
    source = datum.Source.from_arrow(batch())

    with pytest.raises(TypeError):
        source.via(datum.Flow.map_add(1))


def test_udf_exception_preserves_python_traceback_text():
    first = batch((1,))
    second = batch((2,))
    calls = 0

    def fail_on_second(incoming):
        nonlocal calls
        calls += 1
        if calls == 2:
            raise ValueError("mid-stream marker")
        return incoming

    stream = datum.Source.from_arrow([first, second]).map_batches(
        fail_on_second,
        output_schema=first.schema,
    )
    assert calls == 0

    with datum.Runtime() as runtime:
        with pytest.raises(datum.StreamError) as raised:
            stream.count_rows(runtime)

    message = str(raised.value)
    assert "Traceback" in message
    assert "ValueError: mid-stream marker" in message


def test_batch_udf_blueprint_can_be_materialized_twice():
    source_batch = batch()
    stream = datum.Source.from_arrow(source_batch).map_batches(
        lambda incoming: incoming,
        output_schema=source_batch.schema,
    )

    with datum.Runtime() as runtime:
        assert stream.count_rows(runtime) == 4
        assert stream.count_rows(runtime) == 4


def test_flow_map_batches_factory_and_batch_sink_terminals():
    source_batch = batch()
    output_schema = one_column_schema()
    flow = datum.Flow.map_batches(
        lambda incoming: pa.record_batch([incoming.column(0)], schema=output_schema),
        output_schema=output_schema,
    )
    source = datum.Source.from_arrow(source_batch).via(flow)

    with datum.Runtime() as runtime:
        assert source.to_mat(datum.BatchSink.count_batches()).run(runtime).wait() == 1
        assert source.to_mat(datum.Sink.arrow_count_rows()).run(runtime).wait() == 4
        table = source.to_mat(datum.Sink.arrow_collect()).run(runtime).wait()

    assert table.to_pydict() == {"id": [1, 2, 3, 4]}


def test_wait_releases_gil_around_udf_stream_execution():
    batches = [batch((idx, idx + 1, idx + 2, idx + 3)) for idx in range(256)]
    stream = datum.Source.from_arrow(batches).map_batches(
        lambda incoming: incoming,
        output_schema=batches[0].schema,
    )

    progress = 0
    running = True

    def spin():
        nonlocal progress
        while running:
            progress += 1
            time.sleep(0)

    worker = threading.Thread(target=spin)
    worker.start()
    try:
        with datum.Runtime() as runtime:
            completion = stream.to_mat(datum.BatchSink.count_rows()).run(runtime)
            before = progress
            assert completion.wait() == 1024
            after = progress
    finally:
        running = False
        worker.join(timeout=5)

    assert after > before


def test_native_rust_arrow_benchmark_helpers():
    result, wall_ns, cpu_ns, rss_kb = native._bench_arrow_rust_operator(
        batch(), "identity", 1
    )
    assert result == 4
    assert wall_ns > 0
    assert cpu_ns >= 0
    assert rss_kb > 0

    compute_result, compute_wall_ns, compute_cpu_ns, compute_rss_kb = (
        native._bench_arrow_rust_operator(batch(), "compute", 1)
    )
    assert compute_result == 4
    assert compute_wall_ns > 0
    assert compute_cpu_ns >= 0
    assert compute_rss_kb > 0


def test_native_python_arrow_benchmark_buckets():
    result = native._bench_arrow_python_udf(
        batch(),
        lambda incoming: incoming,
        batch().schema,
        1,
    )

    rows, wall_ns, _cpu_ns, rss_kb, batches, gil_ns, export_ns, call_ns, import_ns, validate_ns = result
    assert rows == 4
    assert wall_ns > 0
    assert rss_kb > 0
    assert batches == 1
    assert gil_ns > 0
    assert export_ns > 0
    assert call_ns > 0
    assert import_ns >= 0
    assert validate_ns > 0


def test_polars_arrow_capsule_interop_if_available():
    pl = pytest.importorskip("polars")

    frame = pl.DataFrame({"id": [1, 2], "offset": [10, 20]})
    with datum.Runtime() as runtime:
        assert datum.Source.from_arrow(frame).count_rows(runtime) == 2


def test_duckdb_arrow_capsule_interop_if_available():
    duckdb = pytest.importorskip("duckdb")

    reader = duckdb.sql("select 1::bigint as id, 10::bigint as offset").arrow()
    with datum.Runtime() as runtime:
        assert datum.Source.from_arrow(reader).count_rows(runtime) == 1


def test_cookbook_arrow_udf_recipe():
    # #region cookbook-python-arrow-udf
    import pytest
    import pyarrow as pa
    import pyarrow.compute as pc

    import datum

    source_batch = pa.record_batch(
        [
            pa.array([1, 2, 3, 4], type=pa.int64()),
            pa.array([10, 20, 30, 40], type=pa.int64()),
        ],
        names=["id", "offset"],
    )
    output_schema = pa.schema([("total", pa.int64())])

    def add_columns(incoming):
        total = pc.add(incoming.column(0), incoming.column(1))
        return pa.record_batch([total], schema=output_schema)

    with datum.Runtime() as runtime:
        result = (
            datum.Source.from_arrow(source_batch)
            .map_batches(add_columns, output_schema=output_schema)
            .run_collect(runtime)
        )

    print(result.to_pydict())
    assert result.to_pydict() == {"total": [11, 22, 33, 44]}

    def fail_cleanly(_incoming):
        raise ValueError("bad batch marker")

    failing = datum.Source.from_arrow(source_batch).map_batches(
        fail_cleanly,
        output_schema=source_batch.schema,
    )
    with datum.Runtime() as runtime:
        with pytest.raises(datum.StreamError, match="ValueError: bad batch marker"):
            failing.count_rows(runtime)
    # #endregion cookbook-python-arrow-udf

Expected output:

text
{'total': [11, 22, 33, 44]}

Watch: map_batches is the vectorized tier. Scalar Python callables are accepted by map / filter / flat_map, but they still run as Arrow UDF batches and keep the per-element loop in Python.

Run a Python UDF remotely over zstd ​

Problem: run a Python Arrow UDF on a remote Datum worker pool, authenticate with a token, and choose a wire format for a constrained link.

Start a connect server with udf_workers, connect with the server token, and run the same map_batches graph remotely. For slow links, prefer arrow-ipc-zstd: the PY-3c wire-format A/B records it as the remote/constrained default, with the best rows/s on both 100 Mbit/s rows (benchmark record).

python
import os
import time

import pytest

import datum
import datum._native as native

pa = pytest.importorskip("pyarrow")
pc = pytest.importorskip("pyarrow.compute")

# The default published wheel is TCP-only; QUIC requires a source build with the
# `quic` feature. Skip QUIC round-trip tests cleanly when it is not compiled in.
quic_only = pytest.mark.skipif(
    not native._connect_quic_supported(),
    reason="datum-stream built without QUIC support (TCP-only wheel)",
)


def batch(values=(1, 2, 3, 4)):
    left = pa.array(values, type=pa.int64())
    right = pa.array([value * 10 for value in values], type=pa.int64())
    return pa.record_batch([left, right], names=["left", "right"])


def large_batch(rows=65_536):
    left = pa.array(range(rows), type=pa.int64())
    right = pa.array(range(rows, rows * 2), type=pa.int64())
    return pa.record_batch([left, right], names=["left", "right"])


def wait_until(predicate, timeout=5.0):
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        if predicate():
            return
        time.sleep(0.01)
    assert predicate()


def connect_pair(udf_workers=1):
    server = datum.connect.serve(udf_workers=udf_workers)
    session = datum.connect.connect(server.addr, token=server.token)
    return server, session


def require_udf_shm_root(tmp_path, server, session):
    roots = list(tmp_path.glob("datum-connect-udf-*"))
    if not roots:
        session.close()
        server.close()
        pytest.skip("datum-py was built without the default-off udf-shm feature")
    assert len(roots) == 1
    return roots[0]


WIRE_FORMATS = ["arrow-ipc", "arrow-ipc-lz4", "arrow-ipc-zstd", "parquet"]


def balance_merge_graph():
    def build(builder):
        balance = builder.add(datum.Balance(2))
        merge = builder.add(datum.Merge(2))
        builder.connect(balance.outlet(0), merge.inlet(0))
        builder.connect(balance.outlet(1), merge.inlet(1))
        return datum.FlowShape(balance.inlet(), merge.outlet())

    return datum.GraphDsl.create(build)


def broadcast_zip_graph():
    def build(builder):
        broadcast = builder.add(datum.Broadcast(2))
        zip_stage = builder.add(datum.Zip())
        builder.connect(broadcast.outlet(0), zip_stage.in0())
        builder.connect(broadcast.outlet(1), zip_stage.in1())
        return datum.FlowShape(broadcast.inlet(), zip_stage.outlet())

    return datum.GraphDsl.create(build)


def typed_partition_merge_graph():
    def build(builder):
        partition = builder.add(
            datum.Partition(2, strategy=datum.PartitionStrategy.MODULO)
        )
        merge = builder.add(datum.Merge(2))
        builder.connect(partition.outlet(0), merge.inlet(0))
        builder.connect(partition.outlet(1), merge.inlet(1))
        return datum.FlowShape(partition.inlet(), merge.outlet())

    return datum.GraphDsl.create(build)


def test_named_kernel_connect_plan_round_trips_and_matches_typed_enum_encoding():
    assert native._debug_named_kernel_connect_plans_match()

    server, session = connect_pair()
    try:
        graph = (
            datum.Source.range(0, 20)
            .via(datum.Flow.map_multiply(3))
            .filter_greater_or_equal(21)
            .take(5)
            .to_mat(datum.Sink.collect())
        )
        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait()

        assert session.run(graph) == expected
        assert session.run(graph) == expected
    finally:
        session.close()
        server.close()


def test_typed_partition_and_fold_product_round_trip_over_connect():
    graph = (
        datum.Source.range(1, 6)
        .via_graph(typed_partition_merge_graph())
        .to_mat(datum.Sink.fold_product())
    )
    with datum.Runtime() as runtime:
        expected = graph.run(runtime).wait()
    assert expected == 120

    server, session = connect_pair()
    try:
        assert session.run(graph) == expected
    finally:
        session.close()
        server.close()


def test_connect_linear_callable_lambdas_match_in_process_collect():
    server, session = connect_pair(udf_workers=2)
    try:
        graph = (
            datum.Source.range(0, 20)
            .map(lambda item: item + 1, batch_size=4)
            .filter(lambda item: item % 3 == 0, batch_size=4)
            .flat_map(lambda item: [item, -item], batch_size=4)
            .take(8)
            .to_mat(datum.Sink.collect())
        )
        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait()

        assert session.run(graph) == expected
    finally:
        session.close()
        server.close()


def test_connect_junction_graph_matches_in_process_result():
    server, session = connect_pair()
    try:
        graph = (
            datum.Source.range(0, 12)
            .via_graph(balance_merge_graph())
            .to_mat(datum.Sink.collect())
        )
        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait()

        assert session.run(graph) == expected
    finally:
        session.close()
        server.close()


def test_connect_broadcast_zip_tuple_flow_matches_in_process():
    server, session = connect_pair()
    try:
        graph = broadcast_zip_graph()
        values = [1, 2, 3, 4]

        expected = graph.run_with_input(values)
        assert session.run_with_input(graph, values) == expected
    finally:
        session.close()
        server.close()


def test_connect_zip_shape_uneven_inputs_matches_in_process():
    server, session = connect_pair()
    try:
        graph = datum.GraphDsl.create(lambda builder: builder.add(datum.Zip()))
        left = [1, 2, 3]
        right = [10]

        expected = graph.run_zip(left, right)
        assert session.run_zip(graph, left, right) == expected
    finally:
        session.close()
        server.close()


def test_connect_concat_ordering_matches_in_process():
    server, session = connect_pair()
    try:
        graph = datum.GraphDsl.create(lambda builder: builder.add(datum.Concat(3)))
        inputs = [[1, 2], [], [3, 4]]

        expected = graph.run_concat(inputs)
        assert session.run_concat(graph, inputs) == expected
    finally:
        session.close()
        server.close()


def test_connect_interleave_segment_ordering_matches_in_process():
    server, session = connect_pair()
    try:
        graph = datum.GraphDsl.create(
            lambda builder: builder.add(datum.Interleave(3, segment_size=2))
        )
        inputs = [[1, 2, 3], [10, 11, 12], [20]]

        expected = graph.run_interleave(inputs)
        assert session.run_interleave(graph, inputs) == expected
    finally:
        session.close()
        server.close()


def test_connect_arrow_batch_pipeline_streams_arrow_results():
    server, session = connect_pair()
    try:
        first = batch((1, 2))
        second = batch((3, 4, 5))
        graph = datum.Source.from_arrow([first, second]).to_mat(datum.BatchSink.collect())

        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait().to_pydict()

        assert session.run(graph).to_pydict() == expected
    finally:
        session.close()
        server.close()


def test_connect_arrow_batch_udfs_match_in_process_results():
    server, session = connect_pair(udf_workers=2)
    try:
        source_batch = batch()
        projection_schema = pa.schema([("left", pa.int64())])
        compute_schema = pa.schema([("total", pa.int64())])

        def identity(incoming):
            return incoming

        def project_left(incoming):
            return incoming.select([0])

        def add_columns(incoming):
            total = pc.add(incoming.column(0), incoming.column(1))
            return pa.record_batch([total], schema=compute_schema)

        scenarios = [
            (identity, source_batch.schema),
            (project_left, projection_schema),
            (add_columns, compute_schema),
        ]

        with datum.Runtime() as runtime:
            for udf, output_schema in scenarios:
                graph = (
                    datum.Source.from_arrow(source_batch)
                    .map_batches(udf, output_schema=output_schema)
                    .to_mat(datum.BatchSink.collect())
                )
                expected = graph.run(runtime).wait().to_pydict()
                assert session.run(graph).to_pydict() == expected
    finally:
        session.close()
        server.close()


def test_connect_python_udf_shared_memory_mmap_round_trip_and_cleanup(
    monkeypatch, tmp_path
):
    source_batch = large_batch()
    graph = (
        datum.Source.from_arrow(source_batch)
        .map_batches(lambda incoming: incoming, output_schema=source_batch.schema)
        .to_mat(datum.BatchSink.collect())
    )
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM", "off")
    pipe_server, pipe_session = connect_pair()
    try:
        pipe_result = pipe_session.run(graph).to_pydict()
    finally:
        pipe_session.close()
        pipe_server.close()

    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM", "on")
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_DIR", str(tmp_path))
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_THRESHOLD_BYTES", "1")
    server, session = connect_pair()
    root = require_udf_shm_root(tmp_path, server, session)
    try:
        assert session.run(graph).to_pydict() == pipe_result
        assert list(root.rglob("*.arrow")) == []
    finally:
        session.close()
        server.close()
    assert list(tmp_path.glob("datum-connect-udf-*")) == []


def test_connect_python_udf_shared_memory_unavailable_falls_back_to_pipe(
    monkeypatch, tmp_path
):
    unavailable = tmp_path / "not-a-directory"
    unavailable.write_text("blocks directory creation", encoding="utf-8")
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM", "on")
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_DIR", str(unavailable))
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_THRESHOLD_BYTES", "1")
    source_batch = large_batch()
    server, session = connect_pair()
    try:
        graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(lambda incoming: incoming, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        assert session.run(graph) == source_batch.num_rows
    finally:
        session.close()
        server.close()


def test_connect_python_udf_shared_output_overflow_retries_inline(
    monkeypatch, tmp_path
):
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM", "on")
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_DIR", str(tmp_path))
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_THRESHOLD_BYTES", "1")
    source_batch = large_batch()
    output_schema = pa.schema(
        [(f"value_{index}", pa.int64()) for index in range(6)]
    )

    def expand(incoming):
        columns = [incoming.column(index % 2) for index in range(6)]
        return pa.record_batch(columns, schema=output_schema)

    server, session = connect_pair()
    root = require_udf_shm_root(tmp_path, server, session)
    try:
        graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(expand, output_schema=output_schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        assert session.run(graph) == source_batch.num_rows
        assert list(root.rglob("*.arrow")) == []
    finally:
        session.close()
        server.close()


def test_connect_python_udf_worker_crash_reclaims_shared_segments(
    monkeypatch, tmp_path
):
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM", "on")
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_DIR", str(tmp_path))
    monkeypatch.setenv("DATUM_CONNECT_UDF_SHM_THRESHOLD_BYTES", "1")
    orphan = tmp_path / "datum-connect-udf-4294967294-1-orphan"
    orphan.mkdir()
    (orphan / "output-1.arrow").write_bytes(b"orphan")
    source_batch = large_batch()
    server, session = connect_pair()
    roots = list(tmp_path.glob("datum-connect-udf-*"))
    if len(roots) == 1 and orphan.exists():
        session.close()
        server.close()
        pytest.skip("datum-py was built without the default-off udf-shm feature")
    assert len(roots) == 1
    root = roots[0]
    assert not orphan.exists()
    try:
        def crash(incoming):
            import os

            os._exit(77)

        crash_graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(crash, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        with pytest.raises(datum.StreamError, match="died while executing batch"):
            session.run(crash_graph)
        assert list(root.rglob("*.arrow")) == []

        recovery_graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(lambda incoming: incoming, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        assert session.run(recovery_graph) == source_batch.num_rows
    finally:
        session.close()
        server.close()


@pytest.mark.parametrize("wire_format", WIRE_FORMATS)
def test_connect_wire_formats_match_arrow_ipc_for_identity_and_compute(wire_format):
    server, session = connect_pair(udf_workers=2)
    try:
        source_batch = batch()
        compute_schema = pa.schema([("total", pa.int64())])

        def identity(incoming):
            return incoming

        def add_columns(incoming):
            total = pc.add(incoming.column(0), incoming.column(1))
            return pa.record_batch([total], schema=compute_schema)

        scenarios = [
            (identity, source_batch.schema),
            (add_columns, compute_schema),
        ]

        with datum.Runtime() as runtime:
            for udf, output_schema in scenarios:
                graph = (
                    datum.Source.from_arrow(source_batch)
                    .map_batches(udf, output_schema=output_schema)
                    .to_mat(datum.BatchSink.collect())
                )
                expected = graph.run(runtime).wait().to_pydict()
                assert session.run(graph, wire_format=wire_format).to_pydict() == expected
    finally:
        session.close()
        server.close()


def test_connect_unavailable_wire_format_rejects_cleanly():
    server = datum.connect.serve(udf_workers=1, wire_formats=["arrow-ipc"])
    try:
        with pytest.raises(datum.StreamError, match="no mutually supported wire format"):
            datum.connect.connect(
                server.addr,
                token=server.token,
                wire_formats=["parquet"],
            )
    finally:
        server.close()


def test_connect_wire_format_fallback_and_unnegotiated_override():
    server = datum.connect.serve(udf_workers=1, wire_formats=["arrow-ipc"])
    session = None
    try:
        session = datum.connect.connect(
            server.addr,
            token=server.token,
            wire_formats=["parquet", "arrow-ipc"],
        )
        graph = datum.Source.from_arrow(batch()).to_mat(datum.BatchSink.count_rows())
        assert session.default_wire_format == "arrow-ipc"
        assert session.negotiated_wire_formats() == ["arrow-ipc"]
        assert session.run(graph) == 4
        with pytest.raises(datum.StreamError, match="not negotiated"):
            session.run(graph, wire_format="parquet")
    finally:
        try:
            if session is not None:
                session.close()
        except Exception:
            pass
        server.close()


def test_connect_python_udf_exception_preserves_traceback_text():
    server, session = connect_pair()
    try:
        source_batch = batch()

        def fail(incoming):
            raise ValueError("remote traceback marker")

        graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(fail, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )

        with pytest.raises(datum.StreamError) as raised:
            session.run(graph)

        message = str(raised.value)
        assert "Traceback" in message
        assert "ValueError: remote traceback marker" in message
    finally:
        session.close()
        server.close()


def test_connect_python_udf_worker_crash_is_contained_and_pool_recovers():
    server, session = connect_pair()
    try:
        source_batch = batch()

        def crash(incoming):
            import os

            os._exit(77)

        crash_graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(crash, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )

        with pytest.raises(datum.StreamError) as raised:
            session.run(crash_graph)
        message = str(raised.value)
        assert "Python UDF worker" in message
        assert "died while executing batch" in message

        recovery_graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(lambda incoming: incoming, output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        assert session.run(recovery_graph) == source_batch.num_rows
    finally:
        session.close()
        server.close()


def test_connect_python_udf_is_reused_across_executions(tmp_path):
    server, session = connect_pair()
    try:
        marker = tmp_path / "unpickle-count.txt"
        source_batch = batch()

        class ReuseProbe:
            def __init__(self, path):
                self.path = str(path)

            def __setstate__(self, state):
                self.__dict__.update(state)
                with open(self.path, "a", encoding="utf-8") as handle:
                    handle.write("load\n")

            def __call__(self, incoming):
                return incoming

        graph = (
            datum.Source.from_arrow(source_batch)
            .map_batches(ReuseProbe(marker), output_schema=source_batch.schema)
            .to_mat(datum.BatchSink.count_rows())
        )

        assert session.run(graph) == source_batch.num_rows
        assert session.run(graph) == source_batch.num_rows
        assert marker.read_text(encoding="utf-8").splitlines() == ["load"]
    finally:
        session.close()
        server.close()


def test_connect_python_udf_cpu_bound_work_scales_with_worker_pool():
    if usable_cpu_count() < 4:
        pytest.skip("parallelism proof needs at least four schedulable CPUs")

    batches = [batch((idx, idx + 1)) for idx in range(0, 16, 2)]

    def cpu_bound(incoming):
        acc = 0x9E3779B9
        for idx in range(1_200_000):
            acc = ((acc ^ idx) * 16_777_619) & 0xFFFFFFFF
        if acc == -1:
            raise AssertionError("unreachable")
        return incoming

    one_worker = timed_remote_count(1, batches, cpu_bound)
    four_workers = timed_remote_count(4, batches, cpu_bound)
    assert one_worker / four_workers > 2.0


def usable_cpu_count():
    if hasattr(os, "sched_getaffinity"):
        return len(os.sched_getaffinity(0))
    return os.cpu_count() or 1


def timed_remote_count(udf_workers, batches, udf):
    server, session = connect_pair(udf_workers=udf_workers)
    try:
        graph = (
            datum.Source.from_arrow(batches)
            .map_batches(udf, output_schema=batches[0].schema)
            .to_mat(datum.BatchSink.count_rows())
        )
        started = time.perf_counter()
        assert session.run(graph) == sum(batch.num_rows for batch in batches)
        return time.perf_counter() - started
    finally:
        session.close()
        server.close()


def test_connect_bad_token_rejected_before_execute():
    server = datum.connect.serve(udf_workers=1)
    try:
        with pytest.raises(datum.StreamError, match="bad token"):
            datum.connect.connect(server.addr, token="wrong")
    finally:
        server.close()


def test_connect_cancel_mid_stream_cleans_execution():
    server, session = connect_pair()
    try:
        batches = [batch((idx, idx + 1)) for idx in range(0, 20, 2)]
        graph = datum.Source.from_arrow(batches).to_mat(datum.BatchSink.collect())

        execution = session.start(graph, initial_demand=1)
        first = execution.next_batch()
        assert first.num_rows == 2
        assert server.active_executions() >= 1

        execution.cancel()
        wait_until(lambda: server.active_executions() == 0)
    finally:
        session.close()
        server.close()


def test_connect_client_disconnect_cleans_server_execution():
    server, session = connect_pair()
    try:
        batches = [batch((idx, idx + 1)) for idx in range(0, 20, 2)]
        graph = datum.Source.from_arrow(batches).to_mat(datum.BatchSink.collect())

        session.start(graph, initial_demand=0)
        wait_until(lambda: server.active_executions() >= 1)
        session.close()
        wait_until(lambda: server.active_executions() == 0)
    finally:
        server.close()


def test_connect_plan_version_mismatch_is_clean_error():
    server, session = connect_pair()
    try:
        graph = datum.Source.from_iter([1, 2, 3]).to_mat(datum.Sink.collect())
        with pytest.raises(datum.StreamError, match="unsupported Datum plan version"):
            session._debug_run_with_plan_version(graph, 999)
    finally:
        session.close()
        server.close()


def test_cookbook_remote_udf_zstd_recipe():
    # #region cookbook-python-remote-zstd
    import pyarrow as pa
    import pyarrow.compute as pc

    import datum

    server = datum.connect.serve(
        udf_workers=4,
        wire_formats=["arrow-ipc-zstd", "arrow-ipc"],
    )
    session = datum.connect.connect(
        server.addr,
        token=server.token,
        wire_formats=["arrow-ipc-zstd", "arrow-ipc"],
    )
    try:
        batch = pa.record_batch(
            [
                pa.array([1, 2, 3, 4], type=pa.int64()),
                pa.array([10, 20, 30, 40], type=pa.int64()),
            ],
            names=["left", "right"],
        )
        output_schema = pa.schema([("total", pa.int64())])

        def add_columns(incoming):
            total = pc.add(incoming.column(0), incoming.column(1))
            return pa.record_batch([total], schema=output_schema)

        graph = (
            datum.Source.from_arrow(batch)
            .map_batches(add_columns, output_schema=output_schema)
            .to_mat(datum.BatchSink.collect())
        )
        result = session.run(graph, wire_format="arrow-ipc-zstd")

        print(result.to_pydict())
        assert session.default_wire_format == "arrow-ipc-zstd"
        assert result.to_pydict() == {"total": [11, 22, 33, 44]}
    finally:
        session.close()
        server.close()
    # #endregion cookbook-python-remote-zstd


def test_connect_new_junction_graph_old_plan_version_rejects_cleanly():
    server, session = connect_pair()
    try:
        graph = broadcast_zip_graph()
        with pytest.raises(datum.StreamError, match="unsupported Datum plan version"):
            session._debug_run_with_input_plan_version(graph, [1, 2, 3], 1)
    finally:
        session.close()
        server.close()


def test_connect_orderly_close_keeps_stderr_clean(capfd):
    server, session = connect_pair()
    try:
        graph = datum.Source.range(0, 4).to_mat(datum.Sink.collect())
        assert session.run(graph) == [0, 1, 2, 3]
        capfd.readouterr()
        session.close()
        server.close()
        captured = capfd.readouterr()
        assert captured.err == ""
    finally:
        try:
            session.close()
        except Exception:
            pass
        server.close()


def quic_connect_pair(udf_workers=1):
    """A serve/connect pair over the encrypted QUIC carrier.

    The server uses its self-signed localhost dev certificate and the client
    skips certificate verification, the localhost dev default.
    """
    server = datum.connect.serve(udf_workers=udf_workers, transport="quic")
    session = datum.connect.connect(
        server.addr, token=server.token, transport="quic"
    )
    return server, session


@quic_only
def test_connect_quic_round_trip_matches_in_process():
    server, session = quic_connect_pair()
    try:
        graph = (
            datum.Source.range(0, 20)
            .via(datum.Flow.map_multiply(3))
            .filter_greater_or_equal(21)
            .take(5)
            .to_mat(datum.Sink.collect())
        )
        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait()

        assert session.run(graph) == expected
        # A second run reuses the same QUIC session/stream.
        assert session.run(graph) == expected
    finally:
        session.close()
        server.close()


@quic_only
def test_connect_quic_arrow_batch_pipeline_streams_arrow_results():
    server, session = quic_connect_pair()
    try:
        first = batch((1, 2))
        second = batch((3, 4, 5))
        graph = datum.Source.from_arrow([first, second]).to_mat(datum.BatchSink.collect())

        with datum.Runtime() as runtime:
            expected = graph.run(runtime).wait().to_pydict()

        assert session.run(graph).to_pydict() == expected
    finally:
        session.close()
        server.close()


@quic_only
def test_connect_quic_start_demand_cancel_terminal_semantics():
    server, session = quic_connect_pair()
    try:
        graph = datum.Source.range(0, 100).to_mat(datum.Sink.collect())
        # Start without draining, then cancel: the Cancel/Ack round-trip must
        # complete cleanly over QUIC exactly as it does over TCP.
        execution = session.start(graph, initial_demand=0)
        execution.cancel()
    finally:
        session.close()
        server.close()


@quic_only
def test_connect_quic_orderly_close_keeps_stderr_clean(capfd):
    server, session = quic_connect_pair()
    try:
        graph = datum.Source.range(0, 4).to_mat(datum.Sink.collect())
        assert session.run(graph) == [0, 1, 2, 3]
        capfd.readouterr()
        session.close()
        server.close()
        captured = capfd.readouterr()
        assert captured.err == ""
    finally:
        try:
            session.close()
        except Exception:
            pass
        server.close()


@quic_only
def test_connect_quic_bad_token_rejected():
    server = datum.connect.serve(udf_workers=1, transport="quic")
    try:
        with pytest.raises(datum.StreamError):
            datum.connect.connect(server.addr, token="wrong", transport="quic")
    finally:
        server.close()


def test_connect_unknown_transport_rejected():
    with pytest.raises(datum.BuildError, match="unknown transport"):
        datum.connect.serve(transport="udp")


def test_connect_tcp_rejects_quic_only_options():
    with pytest.raises(datum.BuildError):
        datum.connect.serve(cert="not-with-tcp")
    with pytest.raises(datum.BuildError):
        datum.connect.connect("127.0.0.1:1", insecure_skip_verify=True)

Expected output:

text
{'total': [11, 22, 33, 44]}

Watch: token auth is the trust boundary, not a sandbox. Remote Python UDFs execute client-supplied code in worker subprocesses for isolation and parallelism; only accept trusted clients.

Event-sourced counter with replies ​

Problem: persist counter changes and reply only after each journal write is confirmed.

Use EventSourcedBehavior::with_enforced_replies so every command branch must return a ReplyEffect. then_reply sees the confirmed post-event state. The serializer here is intentionally small and is reused by the next persistence recipes.

rust
use std::{sync::Arc, time::Duration};

use datum::{ActorFlow, ReplyPort, Sink, Source};
use datum_persistence::{
    Effect, EventSourcedBehavior, EventSourcedMessage, EventSourcedRuntime, IdentityAdapter,
    MemoryJournal, PersistenceError, PersistenceId, PersistenceResult, PersistenceSerializer,
    SerializerCatalog, SerializerId,
};

#[derive(Clone, Copy)]
struct I64Serializer;

impl PersistenceSerializer<i64> for I64Serializer {
    fn serializer_id(&self) -> SerializerId {
        SerializerId::new(41)
    }

    fn manifest(&self, _value: &i64) -> String {
        "counter-i64-v1".into()
    }

    fn serialize(&self, value: &i64) -> PersistenceResult<Vec<u8>> {
        Ok(value.to_be_bytes().to_vec())
    }

    fn deserialize(&self, manifest: &str, payload: &[u8]) -> PersistenceResult<i64> {
        if manifest != "counter-i64-v1" || payload.len() != 8 {
            return Err(PersistenceError::Serialization("invalid counter frame".into()));
        }
        let mut bytes = [0_u8; 8];
        bytes.copy_from_slice(payload);
        Ok(i64::from_be_bytes(bytes))
    }
}

fn counter_catalog() -> SerializerCatalog<i64> {
    SerializerCatalog::new(I64Serializer)
}

enum CounterCommand {
    Add { amount: i64, reply_to: ReplyPort<i64> },
}

let behavior = EventSourcedBehavior::with_enforced_replies(
    PersistenceId::new("counter", "counter-1").unwrap(),
    0_i64,
    |_, command| {
        let CounterCommand::Add { amount, reply_to } = command;
        Effect::persist(amount).then_reply(reply_to, |state| *state)
    },
    |state: &i64, event: &i64| state + event,
)
.with_event_adapter(IdentityAdapter, counter_catalog());
let handle = behavior
    .spawn(EventSourcedRuntime::new(Arc::new(MemoryJournal::new())))
    .await
    .unwrap();

let completion = Source::from_iter([2_i64, 3])
    .via(ActorFlow::ask(
        handle.actor_ref().clone(),
        1,
        Duration::from_secs(2),
        |amount, reply_to| {
            EventSourcedMessage::command(CounterCommand::Add { amount, reply_to })
        },
    ))
    .run_with(Sink::collect())
    .unwrap();
let replies = tokio::task::spawn_blocking(move || completion.wait())
    .await
    .unwrap()
    .unwrap();
assert_eq!(replies, vec![2, 5]);
handle.stop();

Watch: use datum::ReplyPort with ActorFlow::ask. Cluster sharding has its own serializable ReplyPort; the sharded recipe below sends that port from then_run after persistence.

Tune snapshots and retention ​

Problem: bound replay time while retaining enough history and snapshots for recovery.

Provide a state adapter/catalog and a snapshot store, then configure count-based retention. This example snapshots every 1,000 events, keeps three snapshot windows, and deletes events older than the retained window only after a snapshot succeeds.

rust
use std::sync::Arc;

use datum_persistence::{
    Effect, EventSourcedBehavior, EventSourcedRuntime, IdentityAdapter, MemoryJournal,
    MemorySnapshotStore, PersistenceId, RetentionCriteria,
};

let retention = RetentionCriteria::snapshot_every(1_000, 3)
    .unwrap()
    .with_delete_events_on_snapshot();
let behavior = EventSourcedBehavior::new(
    PersistenceId::new("counter", "snapshot-counter").unwrap(),
    0_i64,
    |_, delta: i64| Effect::persist(delta),
    |state: &i64, event: &i64| state + event,
)
.with_event_adapter(IdentityAdapter, counter_catalog())
.with_state_adapter(IdentityAdapter, counter_catalog())
.with_retention(retention);

let journal = Arc::new(MemoryJournal::new());
let snapshots = Arc::new(MemorySnapshotStore::new());
let handle = behavior
    .spawn(
        EventSourcedRuntime::new(journal)
            .with_snapshot_store(snapshots),
    )
    .await
    .unwrap();

Watch: snapshot_every rejects zero frequency or retention count. Snapshot cleanup is asynchronous; observe DeleteEvents* and DeleteSnapshots* signals if shutdown must wait for a retention cycle. Omitting with_delete_events_on_snapshot keeps journal history while still pruning old snapshots.

Project events by tag into a read model ​

Problem: build a bounded read model from all current events carrying a tag.

The managed projection runtime now reopens the query from its durable cursor, batches ALO cursor saves, filters resume-boundary duplicates, applies handler recovery and restart backoff, exposes pause/offset management, and performs an acknowledged final flush. The application supplies a restartable source provider and an idempotent handler.

Memory and RocksDB tag queries use strict Offset::Sequence positions. This small provider turns that query capability into the fresh-query factory required by the runtime. Keep query_identity stable and scoped to the backend/database; changing its manifest is an explicit cursor migration.

rust
use std::{cmp::Ordering, sync::Arc};

use datum_persistence::{
    Tag,
    projection::{
        CursorPosition, ProjectionCursor, ProjectionError, ProjectionResult, ProjectionSource,
        ProviderManifest, SourceProvider,
    },
    query::{EventEnvelope, EventsByTagQuery, Offset, QueryMode},
};
use futures_util::StreamExt;

/// Source provider for `events_by_tag` implementations that use strict
/// `Offset::Sequence` resume positions, such as memory and RocksDB.
#[derive(Clone)]
struct SequenceEventsByTagProvider {
    query: Arc<dyn EventsByTagQuery>,
    tag: Tag,
    manifest: ProviderManifest,
}

impl SequenceEventsByTagProvider {
    fn new(
        query: Arc<dyn EventsByTagQuery>,
        query_identity: &str,
        tag: Tag,
    ) -> ProjectionResult<Self> {
        let manifest = ProviderManifest::new(format!(
            "sequence-events-by-tag-v1|query={query_identity}|tag={tag}"
        ))?;
        Ok(Self {
            query,
            tag,
            manifest,
        })
    }
}

#[async_trait::async_trait]
impl SourceProvider<EventEnvelope> for SequenceEventsByTagProvider {
    fn manifest(&self) -> &ProviderManifest {
        &self.manifest
    }

    async fn source_from(
        &self,
        cursor: Option<&ProjectionCursor>,
        mode: QueryMode,
    ) -> ProjectionResult<ProjectionSource<EventEnvelope>> {
        let offset = match cursor {
            Some(cursor) => {
                cursor.ensure_manifest(&self.manifest)?;
                Offset::Sequence(decode_sequence(cursor.position())?)
            }
            None => Offset::NoOffset,
        };
        let stream = self
            .query
            .events_by_tag(&self.tag, offset, mode)
            .await
            .map_err(|error| ProjectionError::Source(error.to_string()))?;
        Ok(Box::pin(stream.map(|item| {
            item.map_err(|error| ProjectionError::Source(error.to_string()))
        })))
    }

    fn extract_cursor(&self, envelope: &EventEnvelope) -> ProjectionResult<ProjectionCursor> {
        let Offset::Sequence(sequence) = envelope.offset else {
            return Err(ProjectionError::Source(
                "sequence tag provider received a non-sequence offset".into(),
            ));
        };
        Ok(ProjectionCursor::from_position(
            self.manifest.clone(),
            format!("sequence-v1|{sequence}"),
        ))
    }

    fn compare_positions(
        &self,
        left: &CursorPosition,
        right: &CursorPosition,
    ) -> ProjectionResult<Ordering> {
        Ok(decode_sequence(left)?.cmp(&decode_sequence(right)?))
    }
}

fn decode_sequence(position: &CursorPosition) -> ProjectionResult<u64> {
    position
        .as_str()
        .strip_prefix("sequence-v1|")
        .and_then(|value| value.parse().ok())
        .ok_or_else(|| {
            ProjectionError::CursorCodec(format!("invalid sequence cursor position {position:?}"))
        })
}

The recipe tags the write-side events as before, but the read side is now a managed ALO projection. The read model keys contributions by (persistence_id, sequence_nr), so replaying an unsaved ALO tail is an idempotent upsert instead of a double increment. It uses the counter_catalog() helper defined in the event-sourced counter recipe above.

rust
use std::{
    collections::BTreeMap,
    sync::{Arc, Mutex},
};

use datum_persistence::{
    Effect, EventSourcedBehavior, EventSourcedRuntime, IdentityAdapter, MemoryJournal,
    PersistenceId, SerializerCatalog, Tag,
    projection::{
        Handler, MemoryProjectionOffsetStore, Projection, ProjectionError, ProjectionId,
        ProjectionOffsetStore, ProjectionResult, ProjectionRuntime, ProjectionStatus,
    },
    query::{EventEnvelope, EventsByTagQuery, QueryMode},
};

#[derive(Default)]
struct CounterReadModel {
    // The event identity makes an at-least-once retry an idempotent upsert.
    deltas: BTreeMap<(String, u64), i64>,
}

struct CounterProjectionHandler {
    model: Arc<Mutex<CounterReadModel>>,
    catalog: SerializerCatalog<i64>,
}

#[async_trait::async_trait]
impl Handler<EventEnvelope> for CounterProjectionHandler {
    async fn process(&mut self, envelope: &EventEnvelope) -> ProjectionResult<()> {
        let delta = self
            .catalog
            .decode(&envelope.payload)
            .map_err(|error| ProjectionError::Handler(error.to_string()))?;
        let event_id = (
            envelope.persistence_id.id().to_owned(),
            envelope.sequence_nr.as_u64(),
        );
        self.model
            .lock()
            .map_err(|_| ProjectionError::Handler("read-model lock poisoned".into()))?
            .deltas
            .entry(event_id)
            .or_insert(delta);
        Ok(())
    }
}

let journal = Arc::new(MemoryJournal::new());
let behavior = EventSourcedBehavior::new(
    PersistenceId::new("counter", "projected-counter").unwrap(),
    0_i64,
    |_, delta: i64| Effect::persist(delta),
    |state: &i64, event: &i64| state + event,
)
.with_event_adapter(IdentityAdapter, counter_catalog())
.with_tagger(|_event| vec![Tag::new("counter")]);
let behavior_handle = behavior
    .spawn(EventSourcedRuntime::new(journal.clone()))
    .await
    .unwrap();
behavior_handle.tell(4).unwrap();
behavior_handle.tell(5).unwrap();
assert_eq!(behavior_handle.state().await.unwrap().state, 9);

let provider = SequenceEventsByTagProvider::new(
    Arc::clone(&journal) as Arc<dyn EventsByTagQuery>,
    "memory:counter-docs",
    Tag::new("counter"),
)
.unwrap();
let read_model = Arc::new(Mutex::new(CounterReadModel::default()));
let handler_model = Arc::clone(&read_model);
let projection_id = ProjectionId::new("counter-totals", "counter");
let projection = Projection::at_least_once(projection_id.clone(), provider, move || {
    CounterProjectionHandler {
        model: Arc::clone(&handler_model),
        catalog: counter_catalog(),
    }
})
.with_query_mode(QueryMode::Current);

let offset_store = Arc::new(MemoryProjectionOffsetStore::new());
let projection_handle = projection
    .spawn(ProjectionRuntime::new(offset_store.clone()))
    .await
    .unwrap();
assert_eq!(
    projection_handle.wait().await.unwrap(),
    ProjectionStatus::Completed
);
assert!(offset_store.load(&projection_id).await.unwrap().is_some());
assert_eq!(read_model.lock().unwrap().deltas.values().sum::<i64>(), 9);
behavior_handle.stop();

Watch: this provider intentionally accepts only sequence offsets; PostgreSQL and Cassandra tag queries need providers for their own cursor shapes. For PostgreSQL scale-out, use the shipped EventsBySlicesSourceProvider. QueryMode::Current flushes and completes successfully; QueryMode::Live stays active until stopped. Delivery guarantees cover emitted envelopes only; see the source-completeness caveats.

Durable-state CAS upsert with conflict handling ​

Problem: update the durable-state SPI directly and recover cleanly when another writer wins the revision race.

Every mutation must use stored_revision + 1. On conflict, read the latest VersionedState, merge or re-evaluate the application update, and retry with the returned expected revision.

rust
use datum_persistence::{
    DurableStateStore, MemoryDurableStateStore, PersistenceError, PersistenceId, Revision,
    SerializedPayload, SerializerId,
};

fn frame(value: u8) -> SerializedPayload {
    SerializedPayload::new(SerializerId::new(7), "counter-state-v1", vec![value])
}

let store = MemoryDurableStateStore::new();
let persistence_id = PersistenceId::new("durable-counter", "counter-1").unwrap();
store
    .upsert(&persistence_id, Revision::new(1), frame(10), None)
    .await
    .unwrap();

let conflict = store
    .upsert(&persistence_id, Revision::new(1), frame(11), None)
    .await;
let expected = match conflict {
    Err(PersistenceError::RevisionConflict {
        expected, actual, ..
    }) => {
        assert_eq!(actual, Revision::new(1));
        expected
    }
    other => panic!("expected revision conflict, got {other:?}"),
};

let latest = store.get(&persistence_id).await.unwrap();
assert_eq!(latest.revision, Revision::new(1));
store
    .upsert(&persistence_id, expected, frame(11), None)
    .await
    .unwrap();
assert_eq!(store.get(&persistence_id).await.unwrap().revision, Revision::new(2));

Watch: do not blindly retry the same revision. A conflict is distinct from PersistenceError::Failure: the conflict proves the attempted CAS did not commit, while a generic failure may have an uncertain outcome.

Sharded event-sourced entity ​

Problem: derive a stable persistence id from each sharded entity identity and recover after passivation or handoff.

Enable the integration feature and add the sharding crate used by the command protocol:

toml
[dependencies]
datum-cluster-sharding = "0.11.3"
datum-persistence = { version = "0.11.3", features = ["sharding"] }

EventSourcedEntity::register replaces the blueprint's placeholder id with type_name|entity_id. Sharding commands use datum_cluster_sharding::ReplyPort, so send the reply from then_run after the event is confirmed.

rust
use std::{sync::Arc, time::Duration};

use datum_cluster_sharding::{
    EntityContext, ReplyPort, ShardingHandle,
    serde::{Deserialize, Serialize},
};
use datum_persistence::{
    Effect, EventSourcedBehavior, EventSourcedRuntime, IdentityAdapter, Journal, PersistenceId,
    sharding::EventSourcedEntity,
};

#[derive(Serialize, Deserialize)]
#[serde(crate = "datum_cluster_sharding::serde")]
enum ShardedCounterCommand {
    Add { amount: i64, reply_to: ReplyPort<i64> },
}

async fn register_counter(
    sharding: &ShardingHandle,
    journal: Arc<dyn Journal>,
) {
    EventSourcedEntity::register(
        sharding,
        "persistent-counter",
        EventSourcedRuntime::new(journal),
        |_context: EntityContext| {
            EventSourcedBehavior::new(
                PersistenceId::from_unique("replaced-by-sharding").unwrap(),
                0_i64,
                |_, command| {
                    let ShardedCounterCommand::Add { amount, reply_to } = command;
                    Effect::persist(amount).then_run(move |state| {
                        let _sent = reply_to.send(*state);
                    })
                },
                |state: &i64, event: &i64| state + event,
            )
            .with_event_adapter(IdentityAdapter, counter_catalog())
        },
    )
    .await
    .unwrap();
}

async fn add(sharding: &ShardingHandle, amount: i64) -> i64 {
    let entity = sharding.entity_ref::<ShardedCounterCommand>(
        "persistent-counter",
        "counter-1",
    );
    entity
        .ask(Duration::from_secs(2), |reply_to| ShardedCounterCommand::Add {
            amount,
            reply_to,
        })
        .await
        .unwrap()
}

Watch: passivation or a behavior stop ends the current actor. The next sharded delivery creates a new incarnation and runs journal recovery before handling that command. A journal failure does not hot-recreate the entity on the failing message; the next delivery is the respawn trigger.

Next steps ​