feat: add NominalAvroWriter primitive for file-only avro output

nominal-streaming/Cargo.toml

@@ -10,6 +10,7 @@ repository = "https://github.com/nominal-io/nominal-streaming"
 [features]
 default = ["logging"]
 logging = ["tracing-subscriber"]
+polars = ["dep:polars"]
 
 
 [dependencies]
@@ -25,6 +26,7 @@ derive_more = { workspace = true }
 futures = { workspace = true }
 nominal-api = { workspace = true }
 parking_lot = { workspace = true }
+polars = { version = "0.52", default-features = false, features = ["dtype-struct", "dtype-array"], optional = true }
 prost = { workspace = true }
 reqwest = { workspace = true }
 serde_json = { workspace = true }

nominal-streaming/src/avro_writer/consumer.rs

@@ -0,0 +1,232 @@
+use std::path::Path;
+use std::sync::atomic::Ordering;
+use std::sync::Arc;
+use std::sync::OnceLock;
+
+use nominal_api::tonic::io::nominal::scout::api::proto;
+use nominal_api::tonic::io::nominal::scout::api::proto::WriteRequestNominal;
+use tracing::error;
+
+use super::error::avro_error_from_consumer_ref;
+use super::error::AvroWriterError;
+use super::stats::PipelineStats;
+use crate::consumer::ConsumerError;
+use crate::consumer::ConsumerResult;
+use crate::consumer::WriteRequestConsumer;
+
+/// Avro file consumer that parallelizes per-series `Record` building across
+/// scoped threads, then feeds the resulting records to a single
+/// `apache_avro::Writer` under a mutex (identical on-disk output to
+/// [`crate::consumer::AvroFileConsumer`]). Reuses
+/// [`crate::consumer::points_to_avro`] for the per-dtype `Value`-tree
+/// construction — only the parallel wrapper is new.
+///
+/// Lives alongside [`super::AvroWriter`] because the parallelism only
+/// benefits the single-file write path — the shared `AvroFileConsumer`
+/// used by network-streaming callers stays untouched.
+pub(super) struct ParallelAvroFileConsumer {
+    writer: Arc<parking_lot::Mutex<apache_avro::Writer<'static, std::fs::File>>>,
+}
+
+impl std::fmt::Debug for ParallelAvroFileConsumer {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ParallelAvroFileConsumer")
+            .finish_non_exhaustive()
+    }
+}
+
+impl ParallelAvroFileConsumer {
+    pub(super) fn new_with_full_path(path: &Path) -> std::io::Result<Self> {
+        std::fs::create_dir_all(path.parent().unwrap_or(path))?;
+        let file = std::fs::OpenOptions::new()
+            .create(true)
+            .truncate(false)
+            .write(true)
+            .open(path)?;
+        let writer = apache_avro::Writer::builder()
+            .schema(&crate::consumer::CORE_AVRO_SCHEMA)
+            .writer(file)
+            .codec(apache_avro::Codec::Snappy)
+            .build();
+        Ok(Self {
+            writer: Arc::new(parking_lot::Mutex::new(writer)),
+        })
+    }
+}
+
+impl WriteRequestConsumer for ParallelAvroFileConsumer {
+    fn consume(&self, request: &WriteRequestNominal) -> ConsumerResult<()> {
+        let records = build_records_parallel(&request.series);
+        self.writer
+            .lock()
+            .extend(records)
+            .map_err(|e| ConsumerError::AvroError(Box::new(e)))?;
+        Ok(())
+    }
+}
+
+/// Build avro `Record`s from a slice of `Series` in parallel using scoped
+/// threads. Each series is independent (its own timestamp/value vectors and
+/// tags); the expensive part — `Value::Union(0, Box::new(Value::Double(x)))`
+/// allocation per point — is embarrassingly parallel.
+///
+/// We use `std::thread::scope` rather than a persistent worker pool so the
+/// dependency graph stays minimal; one OS-thread spawn per chunk is
+/// ~20-50µs, well under 1% of wall at observed call rates.
+fn build_records_parallel(series: &[proto::Series]) -> Vec<apache_avro::types::Record<'static>> {
+    if series.is_empty() {
+        return Vec::new();
+    }
+
+    let num_threads = std::thread::available_parallelism()
+        .map(|n| n.get())
+        .unwrap_or(4)
+        .min(series.len())
+        .max(1);
+    let chunk_size = series.len().div_ceil(num_threads);
+
+    std::thread::scope(|s| {
+        let handles: Vec<_> = series
+            .chunks(chunk_size)
+            .map(|chunk| s.spawn(move || build_records_for_chunk(chunk)))
+            .collect();
+
+        let mut out: Vec<apache_avro::types::Record<'static>> = Vec::with_capacity(series.len());
+        for h in handles {
+            out.extend(h.join().expect("worker panicked"));
+        }
+        out
+    })
+}
+
+fn build_records_for_chunk(chunk: &[proto::Series]) -> Vec<apache_avro::types::Record<'static>> {
+    use apache_avro::types::Record;
+    use apache_avro::types::Value;
+    chunk
+        .iter()
+        .map(|series| {
+            let (timestamps, values) = crate::consumer::points_to_avro(series.points.as_ref());
+
+            let mut record = Record::new(&crate::consumer::CORE_AVRO_SCHEMA)
+                .expect("Failed to create Avro record");
+            record.put(
+                "channel",
+                series
+                    .channel
+                    .as_ref()
+                    .map(|c| c.name.clone())
+                    .unwrap_or_else(|| "values".to_string()),
+            );
+            record.put("timestamps", Value::Array(timestamps));
+            record.put("values", Value::Array(values));
+            record.put("tags", series.tags.clone());
+            record
+        })
+        .collect()
+}
+
+/// Wraps another [`WriteRequestConsumer`] and latches the first error seen,
+/// so [`super::AvroWriter`] can surface disk / avro failures from subsequent
+/// `write` / `flush` / `sync` / `close` calls.
+///
+/// Also records per-call wall time into `stats.consumer_consume_ns` /
+/// `consumer_consume_calls` — the consumer runs on the stream's dispatcher
+/// thread, so its accumulated wall time attributes downstream work to the
+/// pipeline.
+pub(super) struct ErrorLatchingConsumer<C: WriteRequestConsumer> {
+    pub(super) inner: C,
+    pub(super) first_error: Arc<OnceLock<Arc<AvroWriterError>>>,
+    pub(super) stats: Arc<PipelineStats>,
+}
+
+impl<C: WriteRequestConsumer> std::fmt::Debug for ErrorLatchingConsumer<C> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ErrorLatchingConsumer")
+            .finish_non_exhaustive()
+    }
+}
+
+impl<C: WriteRequestConsumer> WriteRequestConsumer for ErrorLatchingConsumer<C> {
+    fn consume(&self, request: &WriteRequestNominal) -> ConsumerResult<()> {
+        let t0 = std::time::Instant::now();
+        let result = self.inner.consume(request);
+        self.stats
+            .consumer_consume_ns
+            .fetch_add(t0.elapsed().as_nanos() as u64, Ordering::Relaxed);
+        self.stats
+            .consumer_consume_calls
+            .fetch_add(1, Ordering::Relaxed);
+        match result {
+            Ok(()) => Ok(()),
+            Err(e) => {
+                // Latch on first occurrence (first-wins via OnceLock::set).
+                // We call `avro_error_from_consumer_ref` rather than consuming
+                // `e` because we must also return `e` to the stream's
+                // dispatcher — `ConsumerError` isn't `Clone`. The helper
+                // reconstructs `Io` from `io::Error::new(kind, msg)` so the
+                // `Io` variant is preserved; avro errors can't be cloned so
+                // they flatten to `Consumer` in the latched copy.
+                if self.first_error.get().is_none() {
+                    let latched = Arc::new(avro_error_from_consumer_ref(&e));
+                    error!("AvroWriter: encoder error latched: {:?}", latched);
+                    let _ = self.first_error.set(latched);
+                }
+                // Return the error to the stream's dispatcher so it logs / backs off.
+                Err(e)
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A consumer that always fails with `PermissionDenied` IoError.
+    /// Used by `error_latch_preserves_io_variant` (unit-struct, fixed kind).
+    #[derive(Debug)]
+    struct PermissionDeniedConsumer;
+
+    impl WriteRequestConsumer for PermissionDeniedConsumer {
+        fn consume(&self, _request: &WriteRequestNominal) -> ConsumerResult<()> {
+            Err(ConsumerError::IoError(std::io::Error::new(
+                std::io::ErrorKind::PermissionDenied,
+                "denied",
+            )))
+        }
+    }
+
+    #[test_log::test]
+    fn error_latch_preserves_io_variant() {
+        let first_error: Arc<OnceLock<Arc<AvroWriterError>>> = Arc::new(OnceLock::new());
+        let consumer = ErrorLatchingConsumer {
+            inner: PermissionDeniedConsumer,
+            first_error: first_error.clone(),
+            stats: Arc::new(PipelineStats::default()),
+        };
+
+        // Construct a minimal WriteRequestNominal to satisfy the trait bound.
+        let req = WriteRequestNominal {
+            series: vec![],
+            session_name: None,
+        };
+
+        // The consume call should fail (pass the error through) and latch.
+        let result = consumer.consume(&req);
+        assert!(result.is_err(), "consumer should have propagated the error");
+
+        // Inspect the latched error.
+        let latched = first_error.get().expect("error should have been latched");
+        match latched.as_ref() {
+            AvroWriterError::Io(io_err) => {
+                assert_eq!(
+                    io_err.kind(),
+                    std::io::ErrorKind::PermissionDenied,
+                    "expected PermissionDenied, got {:?}",
+                    io_err.kind()
+                );
+            }
+            other => panic!("expected AvroWriterError::Io, got {other:?}"),
+        }
+    }
+}

nominal-streaming/src/avro_writer/error.rs

@@ -0,0 +1,59 @@
+use crate::consumer::ConsumerError;
+
+/// Errors produced by [`super::AvroWriter`].
+///
+/// Latched errors surface via `Consumer` for any variant that couldn't be
+/// preserved across the `&ConsumerError` → `AvroWriterError` boundary (avro
+/// errors specifically, because `apache_avro::Error` isn't `Clone`).
+///
+/// These errors are wrapped in `Arc` for idiomatic sharing of a sticky,
+/// cached error across producer threads and repeated calls — once a writer
+/// latches a failure, every subsequent method call returns the same
+/// `Arc<AvroWriterError>` without re-walking the failure path.
+#[derive(Debug, thiserror::Error)]
+pub enum AvroWriterError {
+    /// I/O error from the underlying file (open, write, sync, etc.).
+    /// Latched I/O errors are reconstructed from `io::ErrorKind` + message
+    /// so the `Io` variant is preserved through the error-latching consumer layer.
+    #[error("I/O error: {0}")]
+    Io(#[from] std::io::Error),
+
+    /// Error from the avro writer (schema mismatch, serialization failure).
+    /// Avro errors latch as `Consumer` because `apache_avro::Error` isn't `Clone`.
+    #[error("avro error: {0}")]
+    Avro(#[from] Box<apache_avro::Error>),
+
+    /// Error from the underlying [`crate::consumer::AvroFileConsumer`] sink,
+    /// or a latched avro error (avro errors flatten here because
+    /// `apache_avro::Error` isn't `Clone`).
+    #[error("consumer error: {0}")]
+    Consumer(String),
+
+    /// `write` was called after `close()` completed. Always means caller
+    /// misuse, not encoder failure.
+    #[error("write attempted after close")]
+    SendAfterClose,
+}
+
+impl From<ConsumerError> for AvroWriterError {
+    fn from(e: ConsumerError) -> Self {
+        match e {
+            ConsumerError::IoError(io) => AvroWriterError::Io(io),
+            ConsumerError::AvroError(avro) => AvroWriterError::Avro(avro),
+            other => AvroWriterError::Consumer(other.to_string()),
+        }
+    }
+}
+
+/// Construct an `AvroWriterError` from a `ConsumerError` that the caller
+/// needs to retain (so we can't consume it). Preserves the `Io` variant
+/// via `io::Error::new(kind, msg)`; avro errors flatten to `Consumer`
+/// because `apache_avro::Error` isn't Clone.
+pub(super) fn avro_error_from_consumer_ref(e: &ConsumerError) -> AvroWriterError {
+    match e {
+        ConsumerError::IoError(io) => {
+            AvroWriterError::Io(std::io::Error::new(io.kind(), io.to_string()))
+        }
+        other => AvroWriterError::Consumer(other.to_string()),
+    }
+}

nominal-streaming/src/avro_writer/helpers.rs

@@ -0,0 +1,71 @@
+use std::io;
+use std::path::Path;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::OnceLock;
+use std::time::Duration;
+
+use super::consumer::ErrorLatchingConsumer;
+use super::consumer::ParallelAvroFileConsumer;
+use super::error::AvroWriterError;
+use super::stats::PipelineStats;
+use crate::consumer::WriteRequestConsumer;
+use crate::stream::NominalDatasetStream;
+use crate::stream::NominalStreamOpts;
+
+/// Derive the numbered path for rotation index `index`.
+/// `out.avro` at index 0 → `out_000.avro`.
+pub(super) fn path_for_index(base: &Path, index: usize) -> PathBuf {
+    let stem = base.file_stem().unwrap_or_default().to_string_lossy();
+    let suffix = base
+        .extension()
+        .map(|e| format!(".{}", e.to_string_lossy()))
+        .unwrap_or_default();
+    base.with_file_name(format!("{stem}_{index:03}{suffix}"))
+}
+
+/// Ensure the parent directory exists and truncate the file at `path`.
+///
+/// Returns `io::Error` on failure so callers can surface the error to
+/// their own error type. (Historically this panicked; library-grade APIs
+/// should not.)
+pub(super) fn ensure_parent_and_truncate(path: &Path) -> io::Result<()> {
+    if let Some(parent) = path.parent() {
+        if !parent.as_os_str().is_empty() {
+            std::fs::create_dir_all(parent)?;
+        }
+    }
+    std::fs::File::create(path)?;
+    Ok(())
+}
+
+/// Build a fresh `NominalDatasetStream` with the given consumer and writer opts.
+pub(super) fn open_stream<C: WriteRequestConsumer + 'static>(
+    consumer: C,
+    max_points_per_batch: usize,
+    max_batch_delay: Duration,
+) -> NominalDatasetStream {
+    let stream_opts = NominalStreamOpts {
+        max_points_per_record: max_points_per_batch,
+        max_request_delay: max_batch_delay,
+        max_buffered_requests: 4,
+        request_dispatcher_tasks: 1,
+        base_api_url: String::new(),
+    };
+    NominalDatasetStream::new_with_consumer(consumer, stream_opts)
+}
+
+/// Open a new `ParallelAvroFileConsumer` + `ErrorLatchingConsumer` for the
+/// given path, sharing the `first_error` latch and `stats` counters.
+pub(super) fn open_error_latching_consumer(
+    path: &Path,
+    first_error: Arc<OnceLock<Arc<AvroWriterError>>>,
+    stats: Arc<PipelineStats>,
+) -> io::Result<impl WriteRequestConsumer> {
+    let avro_consumer = ParallelAvroFileConsumer::new_with_full_path(path)?;
+    Ok(ErrorLatchingConsumer {
+        inner: avro_consumer,
+        first_error,
+        stats,
+    })
+}