spiral-sdk 0.1.11

//! Fixed scan metadata and lazy execution lanes sharing Arrow batches.

use std::collections::BTreeMap;
use std::ops::Range;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};

use arrow_array::{RecordBatch, RecordBatchReader};
use arrow_schema::{Schema, SchemaRef};
use async_trait::async_trait;
use futures::task::{ArcWake, AtomicWaker, waker_ref};
use futures::{Stream, StreamExt as _, TryStreamExt as _, stream};
use parking_lot::Mutex;
use tokio::sync::mpsc;
use tokio::task::{JoinHandle, spawn_blocking};
use uuid::Uuid;

use crate::scandal::TableBinding;
pub use crate::scandal::api::{ColumnStatistics, ScanStatistics};
use crate::scandal::expression::Expression;
use crate::scandal::types::require_exact_filter;
pub use crate::scandal::types::{
    FilterClause, FilterSupport, ScanMetricsSnapshot, ScanOptions, ScanRowRange, StorageCounters,
    StorageScope, TableSplit, WindowCounters,
};
use crate::{
    Error, ExecutionDiagnostics, ExecutionOptions, PlanDescription, ResourceAddress, SortColumn,
    TableIdentity, WriteLayout,
};

type MetricsProvider = Arc<dyn Fn() -> Result<Option<ScanMetricsSnapshot>, Error> + Send + Sync>;

/// All lanes of one scan execution share these counters.
/// This handle retains the final counters after stream cancellation.
#[derive(Clone, Default)]
pub struct ScanMetrics {
    provider: Option<Arc<Mutex<Option<MetricsProvider>>>>,
}

impl ScanMetrics {
    /// Adapt a backend's counters without extra work on each row or batch.
    pub(crate) fn from_snapshot_fn(
        provider: impl Fn() -> Result<ScanMetricsSnapshot, Error> + Send + Sync + 'static,
    ) -> Self {
        Self {
            provider: Some(Arc::new(Mutex::new(Some(Arc::new(move || {
                provider().map(Some)
            }))))),
        }
    }

    /// This method returns current counters, or None before execution starts or when the backend
    /// supplies none. It does not read storage. Concurrent updates can span the snapshot.
    pub fn snapshot(&self) -> Result<Option<ScanMetricsSnapshot>, Error> {
        let provider = self
            .provider
            .as_ref()
            .and_then(|provider| provider.lock().clone());
        provider.map_or(Ok(None), |provider| provider())
    }

    fn pending() -> Self {
        Self {
            provider: Some(Arc::new(Mutex::new(None))),
        }
    }

    fn bind(&self, source: Self) {
        if let (Some(provider), Some(_)) = (&self.provider, &source.provider) {
            *provider.lock() = Some(Arc::new(move || source.snapshot()));
        }
    }
}

/// One split of a scan's plan: stored rows of one fragment of the exact version, in stored order.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ScanSplit {
    /// The fragment's position in the version.
    pub fragment: u64,
    /// Half-open stored positions within the fragment, before filtering and deletion removal.
    pub rows: Range<u64>,
    /// Rows the split holds, before filtering, when planning estimates them.
    pub estimated_rows: Option<u64>,
    /// Bytes the split reads, when planning estimates them.
    pub estimated_bytes: Option<u64>,
    /// The stable ID of the fragment the split reads, which change plans and column writes name.
    /// `None` when the scandal does not report it.
    pub fragment_id: Option<Uuid>,
}

/// Stored rows of one split, read in stored order: a unit of work for one consumer of a scan,
/// such as one worker of a distributed job. The application decides which consumer reads which
/// items and in what order.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct WorkItem {
    /// Index of the split in [`Scan::splits`].
    pub split: usize,
    /// Half-open stored positions within the split's rows.
    pub rows: Range<u64>,
}

/// Rows of one split of a scan, by position, for [`Scan::take`].
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Take {
    /// Index of the split in [`Scan::splits`].
    pub split: usize,
    /// Rows of the split by position. Position `p` is the `p`-th row, counted from zero in stored
    /// order, that the scanned version keeps in the split, so a split holds as many positions as
    /// [`Scan::count_rows`] counts for it whole. Positions may come in any order and repeat.
    pub positions: Vec<u64>,
}

/// Projected Arrow schema and clause-by-clause predicate support.
#[derive(Clone, Debug)]
pub(crate) struct ScanInfo {
    /// Projected Arrow schema, including field metadata.
    pub schema: Schema,
    pub source_fields: Option<BTreeMap<String, Uuid>>,
    /// Support for each original top-level conjunction term.
    pub filter_clauses: Vec<FilterClause>,
    pub write_layout: Option<WriteLayout>,
    pub statistics: Option<ScanStatistics>,
}

/// A filter that narrows a scan's splits as they open, such as a bound that tightens while the
/// rows already read are consumed. Each split reads with the filter as it stands when the split
/// opens and may skip every row the filter then rejects, so it must reject only rows the reader
/// does not need.
pub trait SplitFilter: Send + Sync {
    /// The filter as it stands now, or `None` while it rejects nothing.
    fn current(&self) -> Option<Expression>;
}

/// An opened execution binding retaining fixed options and authorized storage.
#[async_trait]
pub(crate) trait ScanSource: Send + Sync {
    /// Start a fresh execution of the whole bound scan, each split narrowed by `narrowing`.
    /// A source may ignore it.
    ///
    /// This operation must not resolve a table name or change the scan's options.
    async fn read(&self, narrowing: Option<Arc<dyn SplitFilter>>) -> Result<BatchStream, Error>;

    /// Fragment count and common ordering, when the source can preserve ordered streams.
    fn ordered_partitions(&self) -> Option<(usize, Vec<SortColumn>)> {
        None
    }

    /// Read the ordered fragments through at most `lanes` merged streams.
    fn read_partitions(
        &self,
        _narrowing: Option<Arc<dyn SplitFilter>>,
        _lanes: usize,
        _batch_rows: usize,
    ) -> Option<Vec<BatchStream>> {
        None
    }

    /// A fresh binding of the scanned Table's storage, which evaluations over the scan's rows
    /// use to read the Table's blobs. Sources without one lend no storage.
    async fn table_binding(&self) -> Result<Option<TableBinding>, Error> {
        Ok(None)
    }

    /// An estimate of the stored rows the bound scan reads. Sources without a plan estimate none.
    fn estimated_rows(&self) -> Option<u64> {
        None
    }

    /// The splits the bound scan reads, in stored order. Sources without a plan have none.
    fn splits(&self) -> &[ScanSplit] {
        &[]
    }

    /// Read the rows of validated `takes` in request order.
    async fn take(&self, _takes: Vec<Take>) -> Result<BatchStream, Error> {
        Err(Error::invalid("this scan has no splits to take rows from"))
    }

    /// Read validated `work` in order.
    async fn read_work(&self, _work: Vec<WorkItem>) -> Result<WorkStream, Error> {
        Err(Error::invalid("this scan has no splits to read"))
    }

    /// The rows each item of validated `work` returns after the scan's filter.
    async fn count_rows(&self, _work: &[WorkItem]) -> Result<Vec<u64>, Error> {
        Err(Error::invalid("this scan has no splits to count"))
    }
}

/// One exact table scan with fixed options, metadata, and execution placement.
#[derive(Clone)]
pub struct Scan {
    pub(crate) address: Option<ResourceAddress>,
    identity: TableIdentity,
    options: ScanOptions,
    schema: SchemaRef,
    source_fields: Option<BTreeMap<String, Uuid>>,
    filter_clauses: Vec<FilterClause>,
    write_layout: Option<WriteLayout>,
    statistics: Option<ScanStatistics>,
    stream_count: usize,
    ordering: Vec<SortColumn>,
    source: Arc<dyn ScanSource>,
}

impl Scan {
    /// Bind validated metadata to a lazy source enforcing these options exactly.
    ///
    /// Unsupported predicates are rejected before reading. Stored-order scans,
    /// limits, and explicit fragment selections use one execution lane.
    pub(crate) fn new(
        identity: TableIdentity,
        options: ScanOptions,
        info: ScanInfo,
        source: Arc<dyn ScanSource>,
    ) -> Result<Self, Error> {
        require_exact_filter(&info.filter_clauses)?;
        let max_streams = options.stream_count()?;
        let ordered = source.ordered_partitions().filter(|(count, _)| {
            (*count == 1
                || (options.partitioned && options.limit.is_none() && options.row_range.is_none()))
                && options.emit.is_none()
        });
        let (stream_count, ordering) = ordered
            .map(|(count, ordering)| (count.min(max_streams), ordering))
            .unwrap_or((max_streams, Vec::new()));
        Ok(Self {
            address: None,
            identity,
            options,
            schema: Arc::new(info.schema),
            source_fields: info.source_fields,
            filter_clauses: info.filter_clauses,
            write_layout: info.write_layout,
            statistics: info.statistics,
            stream_count,
            ordering,
            source,
        })
    }

    /// A fresh binding of the scanned Table's storage, for evaluations that read its blobs.
    pub(crate) async fn table_binding(&self) -> Result<Option<TableBinding>, Error> {
        self.source.table_binding().await
    }

    /// Immutable version read by this scan.
    pub fn identity(&self) -> TableIdentity {
        self.identity
    }

    /// Canonical exact-version address inherited from a path-resolved Table.
    /// ID-only scans expose their provenance through [`Self::identity`].
    pub fn resolved_address(&self) -> Option<&ResourceAddress> {
        self.address.as_ref()
    }

    /// Options fixed when the scan was opened.
    pub fn options(&self) -> &ScanOptions {
        &self.options
    }

    /// Projected schema, available without reading batches or performing I/O.
    pub fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    /// This method returns stable IDs for top-level stored columns, independent of the output
    /// projection. `None` marks scans without stable identities, such as change scans or scans
    /// from older hosts.
    pub fn source_fields(&self) -> Option<&BTreeMap<String, Uuid>> {
        self.source_fields.as_ref()
    }

    /// Preferred storage partitioning and clustering for writes to this exact table.
    pub fn write_layout(&self) -> Option<&WriteLayout> {
        self.write_layout.as_ref()
    }

    /// Stored statistics for this scan, with explicit exactness for counts and bounds.
    pub fn statistics(&self) -> Option<&ScanStatistics> {
        self.statistics.as_ref()
    }

    /// Ordering guaranteed independently within each execution stream, with no global order.
    pub fn ordering(&self) -> &[SortColumn] {
        &self.ordering
    }

    /// Exact support for the accepted predicate clauses.
    pub fn filter_clauses(&self) -> &[FilterClause] {
        &self.filter_clauses
    }

    /// An estimate of the stored rows this scan reads, known without reading any rows.
    ///
    /// It counts every row of the splits planning kept, so the scan's filter, limit, or emit can
    /// return fewer rows. `None` when the plan does not estimate every split.
    pub fn estimated_rows(&self) -> Option<u64> {
        self.source.estimated_rows()
    }

    /// Number of execution lanes available through [`Self::into_streams`].
    pub fn stream_count(&self) -> usize {
        self.stream_count
    }

    /// The splits planning kept, in stored order: the units of work that [`Self::read_work`]
    /// reads.
    pub fn splits(&self) -> &[ScanSplit] {
        self.source.splits()
    }

    /// Read `work` in order, each item's rows in stored order.
    ///
    /// Each batch carries the index of its item in `work` and the stored row after the rows it
    /// covered, so a reader resumes an item from that row without reading any row twice. Items
    /// name splits of this scan by index and stored rows within them. Scans with a limit or an
    /// emit read no work items.
    pub async fn read_work(&self, work: Vec<WorkItem>) -> Result<WorkStream, Error> {
        self.validate_work(&work)?;
        self.source.read_work(work).await
    }

    /// The exact rows each item of `work` returns: stored rows that pass the filter and are not
    /// deleted. Counting reads only the filter's columns, and none without a filter.
    pub async fn count_rows(&self, work: &[WorkItem]) -> Result<Vec<u64>, Error> {
        self.validate_work(work)?;
        self.source.count_rows(work).await
    }

    /// Read rows of this scan's splits by position, in request order: each take's rows in the
    /// order of its positions, a repeated position's row again, the takes one after another.
    ///
    /// This is the read a sampler calls: plan the scan once, count each split's rows with
    /// [`Self::count_rows`], then take the positions it draws. Each split's rows are read in one
    /// pass, so positions that share storage share its reads, and takes beyond one request's
    /// bounds are sent as several. A failed request is resumed from the rows it delivered.
    /// A position past its split fails the call before any rows are delivered. Scans with a
    /// filter, a limit or an emit take no rows.
    pub async fn take(&self, takes: Vec<Take>) -> Result<BatchStream, Error> {
        if self.options.filter.is_some()
            || self.options.limit.is_some()
            || self.options.emit.is_some()
        {
            return Err(Error::invalid(
                "scans with a filter, a limit or an emit take no rows by position",
            ));
        }
        let splits = self.splits().len();
        if let Some(take) = takes.iter().find(|take| take.split >= splits) {
            return Err(Error::invalid(format!(
                "a take names split {} of a scan with {splits} splits",
                take.split
            )));
        }
        self.source.take(takes).await
    }

    fn validate_work(&self, work: &[WorkItem]) -> Result<(), Error> {
        if self.options.limit.is_some() || self.options.emit.is_some() {
            return Err(Error::invalid(
                "scans with a limit or an emit read no work items",
            ));
        }
        let splits = self.splits();
        for item in work {
            let split = splits.get(item.split).ok_or_else(|| {
                Error::invalid(format!(
                    "work names split {} of a scan with {} splits",
                    item.split,
                    splits.len()
                ))
            })?;
            if item.rows.start > item.rows.end
                || item.rows.start < split.rows.start
                || item.rows.end > split.rows.end
            {
                return Err(Error::invalid(format!(
                    "work rows {:?} are outside split {} rows {:?}",
                    item.rows, item.split, split.rows
                )));
            }
        }
        Ok(())
    }

    /// Consume the scan as one batch stream. Reading starts on first poll.
    /// Unordered options permit unordered output even when using one consumer.
    pub fn into_stream(self) -> BatchStream {
        if !self.ordering.is_empty() {
            return self
                .source
                .read_partitions(None, 1, self.options.batch_rows.unwrap_or(65_536))
                .expect("ordered source supplies fixed partitions")
                .pop()
                .expect("one ordered lane");
        }
        self.open_stream(None)
    }

    /// Consume the scan as disjoint streams that together contain its selected rows.
    ///
    /// Ordered lanes merge fixed groups of fragments when [`Self::ordering`] is nonempty.
    /// Other lanes share work dynamically. Consume all lanes concurrently for throughput.
    /// A global limit or explicit split yields one lane. Clone the scan, not a lane,
    /// to start a fresh execution without resolving metadata again.
    /// Dropping the last active lane cancels remaining work, including idle lanes.
    pub fn into_streams(self) -> Vec<ScanStream> {
        self.lanes(None)
    }

    /// [`Self::into_streams`], with each split narrowed by `filter` as it stands when the split
    /// opens.
    pub fn into_narrowed_streams(self, filter: Arc<dyn SplitFilter>) -> Vec<ScanStream> {
        self.lanes(Some(filter))
    }

    fn lanes(self, narrowing: Option<Arc<dyn SplitFilter>>) -> Vec<ScanStream> {
        if !self.ordering.is_empty() {
            return self
                .source
                .read_partitions(
                    narrowing,
                    self.stream_count,
                    self.options.batch_rows.unwrap_or(65_536),
                )
                .expect("ordered source supplies fixed partitions")
                .into_iter()
                .map(|stream| ScanStream {
                    schema: stream.schema(),
                    metrics: stream.metrics(),
                    diagnostics: stream.diagnostics(),
                    stream: Box::pin(stream),
                })
                .collect();
        }
        let count = self.stream_count;
        let schema = self.schema();
        let stream = self.open_stream(narrowing);
        let metrics = stream.metrics();
        let diagnostics = stream.diagnostics();
        let shared = Arc::new(LaneSource {
            state: Mutex::new(LaneState {
                stream: Some(stream),
                active: 0,
                failure: None,
            }),
            waiters: (0..count).map(|_| AtomicWaker::new()).collect(),
        });
        (0..count)
            .map(|index| ScanStream {
                schema: Arc::clone(&schema),
                metrics: metrics.clone(),
                diagnostics: diagnostics.clone(),
                stream: Box::pin(Lane {
                    shared: Arc::clone(&shared),
                    index,
                    active: false,
                }),
            })
            .collect()
    }

    fn open_stream(self, narrowing: Option<Arc<dyn SplitFilter>>) -> BatchStream {
        let schema = self.schema();
        let metrics = if self.options.collect_metrics {
            ScanMetrics::pending()
        } else {
            ScanMetrics::default()
        };
        let execution_metrics = metrics.clone();
        let diagnostics = ExecutionDiagnostics::default();
        let execution_diagnostics = diagnostics.clone();
        let opened = stream::once(async move {
            let batches = self.source.read(narrowing).await?;
            execution_metrics.bind(batches.metrics());
            execution_diagnostics.bind(batches.diagnostics());
            if PlanDescription::row_schema(&batches.schema())
                != PlanDescription::row_schema(&self.schema)
            {
                return Err(Error::arrow(
                    "opened stream schema differs from scan schema",
                ));
            }
            Ok(batches)
        });
        BatchStream::from_stream(schema, opened.try_flatten())
            .with_metrics(metrics)
            .with_diagnostics(diagnostics)
    }
}

/// One execution lane of a fixed scan.
pub struct ScanStream {
    diagnostics: ExecutionDiagnostics,
    schema: SchemaRef,
    metrics: ScanMetrics,
    stream: Pin<Box<dyn Stream<Item = Result<RecordBatch, Error>> + Send>>,
}

impl ScanStream {
    /// Return the counters shared by every lane of this execution.
    pub fn metrics(&self) -> ScanMetrics {
        self.metrics.clone()
    }

    /// Fixed projected schema shared by all lanes of the scan.
    pub fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    /// Read this lane lazily. Dropping the last active lane cancels remaining work.
    pub fn into_stream(self) -> BatchStream {
        BatchStream::from_stream(self.schema, self.stream)
            .with_metrics(self.metrics)
            .with_diagnostics(self.diagnostics)
    }
}

struct LaneSource {
    state: Mutex<LaneState>,
    waiters: Vec<AtomicWaker>,
}

struct LaneState {
    stream: Option<BatchStream>,
    active: usize,
    failure: Option<String>,
}

impl ArcWake for LaneSource {
    fn wake_by_ref(shared: &Arc<Self>) {
        for waiter in &shared.waiters {
            waiter.wake();
        }
    }
}

struct Lane {
    shared: Arc<LaneSource>,
    index: usize,
    active: bool,
}

impl Stream for Lane {
    type Item = Result<RecordBatch, Error>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if !self.active {
            self.shared.state.lock().active += 1;
            self.active = true;
        }
        self.shared.waiters[self.index].register(cx.waker());
        let waker = waker_ref(&self.shared);
        let mut source_context = Context::from_waker(&waker);
        let result = {
            let mut state = self.shared.state.lock();
            if let Some(failure) = &state.failure {
                return Poll::Ready(Some(Err(Error::arrow(failure))));
            }
            let Some(stream) = &mut state.stream else {
                return Poll::Ready(None);
            };
            let result = Pin::new(stream).poll_next(&mut source_context);
            match &result {
                Poll::Ready(Some(Err(error))) => {
                    state.failure = Some(error.to_string());
                    state.stream = None;
                }
                Poll::Ready(None) => state.stream = None,
                _ => {}
            }
            result
        };
        if result.is_ready() {
            LaneSource::wake_by_ref(&self.shared);
        }
        result
    }
}

impl Drop for Lane {
    fn drop(&mut self) {
        self.shared.waiters[self.index].take();
        if self.active {
            let mut state = self.shared.state.lock();
            state.active -= 1;
            if state.active == 0 && state.stream.take().is_some() {
                state.failure = Some("scan execution cancelled".into());
            }
        }
    }
}

/// A batch of one work item and where the item's read stands after it.
#[derive(Clone, Debug)]
pub struct WorkBatch {
    pub batch: RecordBatch,
    /// The batch's item, by index in the work that was read.
    pub item: usize,
    /// The stored row after the rows this batch covered; the item resumes from here.
    pub next_row: u64,
}

/// Batches of work items, in the items' order. A failure terminates the stream.
pub struct WorkStream {
    schema: SchemaRef,
    stream: Pin<Box<dyn Stream<Item = Result<WorkBatch, Error>> + Send>>,
}

impl WorkStream {
    pub(crate) fn new(
        schema: SchemaRef,
        stream: impl Stream<Item = Result<WorkBatch, Error>> + Send + 'static,
    ) -> Self {
        Self {
            schema,
            stream: Box::pin(stream),
        }
    }

    /// Fixed projected schema.
    pub fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    /// Read the next batch asynchronously.
    pub async fn next_batch(&mut self) -> Option<Result<WorkBatch, Error>> {
        self.next().await
    }
}

impl Stream for WorkStream {
    type Item = Result<WorkBatch, Error>;
    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.stream.as_mut().poll_next(context)
    }
}

/// Pull-driven Arrow batches with a fixed schema. A failure terminates the stream.
pub struct BatchStream {
    schema: SchemaRef,
    diagnostics: ExecutionDiagnostics,
    metrics: ScanMetrics,
    stream: Pin<Box<dyn Stream<Item = Result<RecordBatch, Error>> + Send>>,
}

impl BatchStream {
    /// Adapt an asynchronous source without eager polling or additional buffering.
    pub fn from_stream(
        schema: SchemaRef,
        stream: impl Stream<Item = Result<RecordBatch, Error>> + Send + 'static,
    ) -> Self {
        let expected = Arc::clone(&schema);
        let stream = stream::unfold(Some(stream.boxed()), move |source| {
            let expected = Arc::clone(&expected);
            async move {
                let mut source = source?;
                let batch = source.next().await?.and_then(|batch| {
                    if PlanDescription::row_schema(&batch.schema())
                        != PlanDescription::row_schema(&expected)
                    {
                        return Err(Error::arrow("batch schema differs from scan schema"));
                    }
                    Ok(batch)
                });
                let source = batch.is_ok().then_some(source);
                Some((batch, source))
            }
        });
        Self {
            schema,
            metrics: ScanMetrics::default(),
            diagnostics: ExecutionDiagnostics::default(),
            stream: Box::pin(stream),
        }
    }

    /// Set generic execution options on this stream's request schema.
    pub fn with_execution_options(mut self, options: ExecutionOptions) -> Result<Self, Error> {
        self.schema = Arc::new(options.with_schema(&self.schema)?);
        Ok(self)
    }

    /// Retain the execution descriptions observed while consuming this stream.
    /// Profiles update on pulls; dropping the stream marks incomplete observations cancelled.
    pub fn diagnostics(&self) -> ExecutionDiagnostics {
        self.diagnostics.clone()
    }

    pub(crate) fn with_diagnostics(mut self, diagnostics: ExecutionDiagnostics) -> Self {
        self.diagnostics = diagnostics;
        self
    }

    /// Attach execution counters supplied by this stream's backend.
    pub fn with_metrics(mut self, metrics: ScanMetrics) -> Self {
        self.metrics = metrics;
        self
    }

    /// Return execution counters without advancing the stream.
    pub fn metrics(&self) -> ScanMetrics {
        self.metrics.clone()
    }

    /// Adapt a synchronous Arrow reader on a blocking worker with bounded buffering.
    #[expect(
        clippy::disallowed_methods,
        reason = "The open SDK adapts blocking Arrow callbacks without a Vortex runtime"
    )]
    pub fn from_reader(mut reader: Box<dyn RecordBatchReader + Send>) -> Self {
        let schema = reader.schema();
        let (sender, receiver) = mpsc::channel(2);
        let worker = spawn_blocking(move || {
            while !sender.is_closed() {
                let Some(batch) = reader.next() else { break };
                let failed = batch.is_err();
                if sender.blocking_send(batch.map_err(Error::arrow)).is_err() || failed {
                    break;
                }
            }
        });
        Self::from_stream(schema, ReaderStream { receiver, worker })
    }

    /// Fixed projected schema.
    pub fn schema(&self) -> SchemaRef {
        Arc::clone(&self.schema)
    }

    /// Read the next batch asynchronously.
    pub async fn next_batch(&mut self) -> Option<Result<RecordBatch, Error>> {
        self.next().await
    }
}

impl Stream for BatchStream {
    type Item = Result<RecordBatch, Error>;
    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        for _ in 0..32 {
            let next = self.stream.as_mut().poll_next(context);
            if let Poll::Ready(Some(Ok(batch))) = &next {
                match self.diagnostics.observe(batch.schema_ref()) {
                    Ok(true) if batch.num_rows() == 0 => continue,
                    Ok(_) => {}
                    Err(error) => return Poll::Ready(Some(Err(error.into()))),
                }
            }
            return next;
        }
        context.waker().wake_by_ref();
        Poll::Pending
    }
}

impl Drop for BatchStream {
    fn drop(&mut self) {
        self.diagnostics.cancelled();
    }
}

struct ReaderStream {
    receiver: mpsc::Receiver<Result<RecordBatch, Error>>,
    worker: JoinHandle<()>,
}

impl Stream for ReaderStream {
    type Item = Result<RecordBatch, Error>;
    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.receiver.poll_recv(context)
    }
}

impl Drop for ReaderStream {
    fn drop(&mut self) {
        self.receiver.close();
        self.worker.abort();
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::{AtomicUsize, Ordering};

    use arrow_array::Int64Array;
    use arrow_schema::{DataType, Field};
    use uuid::Uuid;

    use super::*;

    struct WorkSource {
        reads: AtomicUsize,
        drops: Arc<AtomicUsize>,
        schema: SchemaRef,
        fail: bool,
    }

    struct WorkLifetime(Arc<AtomicUsize>);
    impl Drop for WorkLifetime {
        fn drop(&mut self) {
            self.0.fetch_add(1, Ordering::SeqCst);
        }
    }

    #[async_trait]
    impl ScanSource for WorkSource {
        async fn read(&self, _: Option<Arc<dyn SplitFilter>>) -> Result<BatchStream, Error> {
            self.reads.fetch_add(1, Ordering::SeqCst);
            let schema = Arc::clone(&self.schema);
            let fail = self.fail;
            let guard = WorkLifetime(Arc::clone(&self.drops));
            let rows = Arc::new(AtomicUsize::new(0));
            let counters = Arc::clone(&rows);
            let metrics = ScanMetrics::from_snapshot_fn(move || {
                Ok(ScanMetricsSnapshot {
                    output_rows: (counters.load(Ordering::Relaxed) as u64).into(),
                    ..Default::default()
                })
            });
            Ok(BatchStream::from_stream(
                Arc::clone(&schema),
                stream::unfold((0, guard), move |(index, guard)| {
                    let schema = Arc::clone(&schema);
                    let rows = Arc::clone(&rows);
                    async move {
                        if index == 100 {
                            return None;
                        }
                        let batch = if fail && index == 1 {
                            Err(Error::arrow("work failed"))
                        } else {
                            Ok(RecordBatch::try_new(
                                schema,
                                vec![Arc::new(Int64Array::from(vec![index]))],
                            )
                            .unwrap())
                        };
                        if batch.is_ok() {
                            rows.fetch_add(1, Ordering::Relaxed);
                        }
                        Some((batch, (index + 1, guard)))
                    }
                }),
            )
            .with_metrics(metrics))
        }
    }

    fn work_scan(fail: bool) -> (Scan, Arc<WorkSource>) {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));
        let source = Arc::new(WorkSource {
            reads: AtomicUsize::new(0),
            drops: Arc::new(AtomicUsize::new(0)),
            schema: Arc::clone(&schema),
            fail,
        });
        let scan = Scan::new(
            TableIdentity {
                resource_id: Uuid::new_v4(),
                version_id: Uuid::new_v4(),
            },
            ScanOptions {
                collect_metrics: true,
                partitioned: true,
                max_streams: Some(2),
                ..Default::default()
            },
            ScanInfo {
                source_fields: None,
                statistics: None,
                write_layout: None,
                schema: schema.as_ref().clone(),
                filter_clauses: vec![],
            },
            source.clone(),
        )
        .unwrap();
        (scan, source)
    }

    #[tokio::test]
    async fn metrics_follow_execution_lifetime_and_replay() {
        let (scan, _) = work_scan(false);
        let mut lanes = scan.clone().into_streams().into_iter();
        let first = lanes.next().unwrap();
        let second = lanes.next().unwrap();
        let first_metrics = first.metrics();
        let second_metrics = second.metrics();
        assert!(first_metrics.snapshot().unwrap().is_none());
        let mut first = first.into_stream();
        let mut second = second.into_stream();
        first.next_batch().await.unwrap().unwrap();
        second.next_batch().await.unwrap().unwrap();
        assert_eq!(first_metrics.snapshot().unwrap().unwrap().output_rows.0, 2);
        assert_eq!(
            first_metrics.snapshot().unwrap(),
            second_metrics.snapshot().unwrap()
        );
        drop(first);
        drop(second);
        assert_eq!(first_metrics.snapshot().unwrap().unwrap().output_rows.0, 2);
        let mut replay = scan.into_stream();
        let replay_metrics = replay.metrics();
        assert!(replay_metrics.snapshot().unwrap().is_none());
        replay.next_batch().await.unwrap().unwrap();
        assert_eq!(replay_metrics.snapshot().unwrap().unwrap().output_rows.0, 1);
        assert_eq!(first_metrics.snapshot().unwrap().unwrap().output_rows.0, 2);
    }

    #[tokio::test]
    async fn lanes_share_work_and_replay_creates_a_fresh_execution() {
        let (scan, source) = work_scan(false);
        let mut lanes = scan
            .clone()
            .into_streams()
            .into_iter()
            .map(ScanStream::into_stream);
        let mut slow = lanes.next().unwrap();
        let mut fast = lanes.next().unwrap();
        assert_eq!(source.reads.load(Ordering::SeqCst), 0);
        let first = slow.next_batch().await.unwrap().unwrap();
        assert_eq!(
            first
                .column(0)
                .as_any()
                .downcast_ref::<Int64Array>()
                .unwrap()
                .value(0),
            0
        );
        for value in 1..100 {
            let batch = fast.next_batch().await.unwrap().unwrap();
            assert_eq!(
                batch
                    .column(0)
                    .as_any()
                    .downcast_ref::<Int64Array>()
                    .unwrap()
                    .value(0),
                value
            );
        }
        assert!(fast.next_batch().await.is_none());
        assert!(slow.next_batch().await.is_none());
        assert_eq!(source.reads.load(Ordering::SeqCst), 1);
        assert_eq!(source.drops.load(Ordering::SeqCst), 1);
        let batches = scan.into_stream().collect::<Vec<_>>().await;
        assert_eq!(batches.len(), 100);
        assert_eq!(source.reads.load(Ordering::SeqCst), 2);
    }

    #[tokio::test]
    async fn lane_failure_and_last_active_lane_cancellation_release_work() {
        let (scan, source) = work_scan(true);
        let mut lanes = scan.into_streams().into_iter().map(ScanStream::into_stream);
        let mut first = lanes.next().unwrap();
        let mut second = lanes.next().unwrap();
        first.next_batch().await.unwrap().unwrap();
        assert!(second.next_batch().await.unwrap().is_err());
        assert!(first.next_batch().await.unwrap().is_err());
        assert!(first.next_batch().await.is_none());
        assert!(second.next_batch().await.is_none());
        assert_eq!(source.drops.load(Ordering::SeqCst), 1);

        let (scan, source) = work_scan(false);
        let mut lanes = scan.into_streams().into_iter().map(ScanStream::into_stream);
        let mut active = lanes.next().unwrap();
        let mut idle = lanes.next().unwrap();
        active.next_batch().await.unwrap().unwrap();
        drop(active);
        assert_eq!(
            source.drops.load(Ordering::SeqCst),
            1,
            "idle lanes must not retain active I/O"
        );
        assert!(idle.next_batch().await.unwrap().is_err());
        assert_eq!(source.reads.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn batch_schema_mismatch_is_terminal() {
        let expected = Arc::new(Schema::new(vec![Field::new(
            "expected",
            DataType::Int64,
            false,
        )]));
        let wrong = RecordBatch::try_new(
            Arc::new(Schema::new(vec![Field::new(
                "wrong",
                DataType::Int64,
                false,
            )])),
            vec![Arc::new(Int64Array::from(vec![1]))],
        )
        .unwrap();
        let valid = RecordBatch::try_new(
            Arc::clone(&expected),
            vec![Arc::new(Int64Array::from(vec![2]))],
        )
        .unwrap();
        let polled = Arc::new(AtomicUsize::new(0));
        let counter = Arc::clone(&polled);
        let source = stream::iter([Ok(wrong), Ok(valid)]).inspect(move |_| {
            counter.fetch_add(1, Ordering::SeqCst);
        });
        let mut batches = BatchStream::from_stream(expected, source);
        assert!(matches!(
            batches.next_batch().await,
            Some(Err(Error::Arrow(_)))
        ));
        assert!(batches.next_batch().await.is_none());
        assert_eq!(polled.load(Ordering::SeqCst), 1);
    }
}