src/sink.rs
//! The Arrow → **skade Iceberg** sink helper. A [`BatchSink`] is what every
//! [`Collector`](super::Collector) streams `RecordBatch`es into; [`run`](super::run)
//! backs it with a real `skade::Table`, while tests back it with an in-memory
//! buffer so the parsers can be asserted as DATA without a warehouse.
//!
//! Buffer-then-flush keeps it simple: collectors `push` batches, [`run`] calls
//! [`BatchSink::flush`] once to land them as one `fast_append` commit (the
//! korp-lite, no-JVM path). For the fixture tests the same batches are read back
//! via [`BatchSink::take_batches`].
//!
//! **Three backings, one sink.** `memory` (tests), `iceberg` (a skade `Table` the
//! leaf opened itself) and — behind `daemon` — [`BatchSink::warehouse`]: an
//! `Arc<dyn WarehouseSink>` the HOST hands in, so a collector's batches land through
//! whatever the host's warehouse seam is (korp: `Streaming::land` on the configured
//! backend row) instead of a skade directory the leaf opened on its own. Every
//! `Collector` streams into the same `&BatchSink`, so the choice of backing is the
//! runner's ([`run`](super::run) vs [`run_with_sink`](super::run_with_sink)), never
//! the collector's.
use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use anyhow::Result;
use arrow_array::RecordBatch;
/// A sink for parsed Arrow batches — memory-backed (tests) or skade-backed (run).
pub struct BatchSink {
buf: Mutex<Vec<RecordBatch>>,
rows: AtomicU64,
batches: AtomicU64,
#[cfg(feature = "skade-sink")]
table: Option<tokio::sync::Mutex<skade::Table>>,
/// The HOST's warehouse seam + the table the batches land in (`daemon`).
#[cfg(feature = "daemon")]
warehouse: Option<(std::sync::Arc<dyn crate::daemon::WarehouseSink>, String)>,
/// Rows the backing reported LANDED (not pushed) — the warehouse's number.
landed: AtomicU64,
}
impl Default for BatchSink {
fn default() -> Self {
Self::memory()
}
}
impl BatchSink {
/// An in-memory sink — batches are buffered and read back with
/// [`take_batches`](Self::take_batches). Used by the fixture tests.
pub fn memory() -> Self {
Self {
buf: Mutex::new(Vec::new()),
rows: AtomicU64::new(0),
batches: AtomicU64::new(0),
#[cfg(feature = "skade-sink")]
table: None,
#[cfg(feature = "daemon")]
warehouse: None,
landed: AtomicU64::new(0),
}
}
/// A **host-warehouse-backed** sink — [`flush`](Self::flush) hands every
/// buffered batch to `sink.append(table, …)` (the host's seam; in korp that is
/// `Streaming::land` on the configured backend row). The host's count is what
/// [`landed`](Self::landed) reports.
#[cfg(feature = "daemon")]
pub fn warehouse(sink: std::sync::Arc<dyn crate::daemon::WarehouseSink>, table: impl Into<String>) -> Self {
Self {
buf: Mutex::new(Vec::new()),
rows: AtomicU64::new(0),
batches: AtomicU64::new(0),
#[cfg(feature = "skade-sink")]
table: None,
warehouse: Some((sink, table.into())),
landed: AtomicU64::new(0),
}
}
/// A skade-backed sink — [`flush`](Self::flush) lands the buffered batches in
/// the Iceberg `table` as one commit.
#[cfg(feature = "skade-sink")]
pub fn iceberg(table: skade::Table) -> Self {
Self {
buf: Mutex::new(Vec::new()),
rows: AtomicU64::new(0),
batches: AtomicU64::new(0),
table: Some(tokio::sync::Mutex::new(table)),
#[cfg(feature = "daemon")]
warehouse: None,
landed: AtomicU64::new(0),
}
}
/// Push one parsed batch. Empty batches are dropped (no-op) so a parser that
/// yields nothing doesn't inflate the batch count.
pub fn push(&self, batch: RecordBatch) -> Result<()> {
if batch.num_rows() == 0 {
return Ok(());
}
self.rows
.fetch_add(batch.num_rows() as u64, Ordering::Relaxed);
self.batches.fetch_add(1, Ordering::Relaxed);
self.buf.lock().unwrap().push(batch);
Ok(())
}
/// Total rows pushed so far.
pub fn rows(&self) -> u64 {
self.rows.load(Ordering::Relaxed)
}
/// Total non-empty batches pushed so far.
pub fn batches(&self) -> u64 {
self.batches.load(Ordering::Relaxed)
}
/// Rows the BACKING reported landed on `flush` — the warehouse's number, never
/// the push tally. `0` for a memory sink and before any flush.
pub fn landed(&self) -> u64 {
self.landed.load(Ordering::Relaxed)
}
/// Land buffered batches in the backing Iceberg table (skade-backed sink) as
/// one `fast_append` commit. No-op for a memory sink. skade recasts each
/// batch to the table schema automatically.
pub async fn flush(&self) -> Result<()> {
#[cfg(feature = "skade-sink")]
if let Some(table) = &self.table {
let drained: Vec<RecordBatch> = std::mem::take(&mut *self.buf.lock().unwrap());
if !drained.is_empty() {
let n: u64 = drained.iter().map(|b| b.num_rows() as u64).sum();
table.lock().await.append(&drained).await?;
self.landed.fetch_add(n, Ordering::Relaxed);
}
}
// The host's seam: `WarehouseSink::append` is BLOCKING (korp's impl drives
// an async `land` on its own runtime), so it runs on the blocking pool — a
// `block_on` inside this async context would panic on a tokio worker.
#[cfg(feature = "daemon")]
if let Some((sink, table)) = &self.warehouse {
let drained: Vec<RecordBatch> = std::mem::take(&mut *self.buf.lock().unwrap());
if !drained.is_empty() {
let sink = std::sync::Arc::clone(sink);
let table = table.clone();
let n = tokio::task::spawn_blocking(move || -> Result<u64> {
let mut n = 0u64;
for b in &drained {
n += sink.append(&table, b)?;
}
Ok(n)
})
.await
.map_err(|e| anyhow::anyhow!("warehouse flush worker died: {e}"))??;
self.landed.fetch_add(n, Ordering::Relaxed);
}
}
Ok(())
}
/// Drain the buffered batches (memory sink) — the fixture tests read the
/// parsed Arrow back here and assert schema + row values.
pub fn take_batches(&self) -> Vec<RecordBatch> {
std::mem::take(&mut *self.buf.lock().unwrap())
}
}
#[cfg(all(test, feature = "daemon"))]
mod warehouse_backing_tests {
use super::*;
use crate::daemon::{MemWarehouse, WarehouseSink};
use arrow_array::{Int64Array, RecordBatch};
use arrow_schema::{DataType, Field, Schema};
use std::sync::Arc;
fn batch(n: i64) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from((0..n).collect::<Vec<_>>()))]).unwrap()
}
/// A sink that swallows every append and reports 0 — the RED arm: a runner that
/// echoed its push tally would say 5; one that reports the SINK's number says 0.
struct Swallow;
impl WarehouseSink for Swallow {
fn append(&self, _t: &str, _b: &RecordBatch) -> anyhow::Result<u64> {
Ok(0)
}
fn rows(&self, _t: &str) -> u64 {
0
}
}
/// `flush` hands every buffered batch to the host's `WarehouseSink::append` under
/// the sink's table, and `landed` is what the HOST reported.
#[tokio::test(flavor = "multi_thread")]
async fn a_warehouse_backed_sink_lands_through_the_hosts_append() {
let mem = MemWarehouse::new();
let sink = BatchSink::warehouse(Arc::new(mem.clone()), "t_host");
sink.push(batch(3)).unwrap();
sink.push(batch(2)).unwrap();
assert_eq!(sink.landed(), 0, "nothing landed before flush");
sink.flush().await.unwrap();
assert_eq!(mem.rows("t_host"), 5, "the host's seam received both batches");
assert_eq!(sink.landed(), 5, "landed is the host's number");
assert_eq!(sink.rows(), 5);
sink.flush().await.unwrap();
assert_eq!(mem.rows("t_host"), 5, "a second flush re-sends nothing");
let swallow = BatchSink::warehouse(Arc::new(Swallow), "t_swallow");
swallow.push(batch(5)).unwrap();
swallow.flush().await.unwrap();
assert_eq!(swallow.rows(), 5, "pushed 5");
assert_eq!(swallow.landed(), 0, "…but the host landed 0, and THAT is what is reported");
}
}