laddu_runtime/

query.rs

use std::sync::Arc;

use laddu_compile::CompiledQuery;
use laddu_data::{
    LadduDataError, LadduDataResult,
    data::{Dataset, EventBatch},
    io::{EventBatchIter, EventSource, ReadPlan, SourceCapabilities, memory::MemorySource},
    schema::Schema,
};
use laddu_expr::{Expr, ValueKind};
use laddu_physics::{
    binning::{BinningAxis, FinalUpperEdge},
    histogram::Histogram,
    joint_histogram::JointHistogram,
};
use num::complex::Complex64;
use serde::{Deserialize, Deserializer, Serialize};

use crate::{Execution, PreparedModel, RuntimeError, RuntimeResult};

/// Comparison operation used by a dataset predicate.
#[derive(Copy, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum Comparison {
    /// Less than.
    Lt,
    /// Less than or equal to.
    Le,
    /// Greater than.
    Gt,
    /// Greater than or equal to.
    Ge,
    /// Equal to.
    Eq,
    /// Not equal to.
    Ne,
}

/// Determines which endpoints are included by an interval predicate.
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub enum IntervalClosure {
    /// Exclude both endpoints.
    Open,
    /// Include only the lower endpoint.
    LeftClosed,
    /// Include only the upper endpoint.
    RightClosed,
    /// Include both endpoints.
    #[default]
    Closed,
}

/// Boolean expression used to select events from a dataset.
#[derive(Clone, Debug)]
pub enum Predicate {
    /// Compare two scalar expressions.
    Compare {
        /// Left-hand expression.
        lhs: Expr,
        /// Comparison operation.
        op: Comparison,
        /// Right-hand expression.
        rhs: Expr,
    },
    /// Require both child predicates to hold.
    And(Box<Self>, Box<Self>),
    /// Require either child predicate to hold.
    Or(Box<Self>, Box<Self>),
    /// Negate a predicate.
    Not(Box<Self>),
    /// Test whether a value lies between two bounds.
    Between {
        /// Expression whose value is tested.
        value: Expr,
        /// Lower bound expression.
        lower: Expr,
        /// Upper bound expression.
        upper: Expr,
        /// Endpoint inclusion policy.
        closure: IntervalClosure,
    },
}

impl Predicate {
    /// Creates a comparison predicate.
    pub fn compare(lhs: impl Into<Expr>, op: Comparison, rhs: impl Into<Expr>) -> Self {
        Self::Compare {
            lhs: lhs.into(),
            op,
            rhs: rhs.into(),
        }
    }

    /// Creates a less-than predicate.
    pub fn lt(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Lt, rhs)
    }
    /// Creates a less-than-or-equal predicate.
    pub fn le(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Le, rhs)
    }
    /// Creates a greater-than predicate.
    pub fn gt(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Gt, rhs)
    }
    /// Creates a greater-than-or-equal predicate.
    pub fn ge(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Ge, rhs)
    }
    /// Creates an equality predicate.
    pub fn eq(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Eq, rhs)
    }
    /// Creates an inequality predicate.
    pub fn ne(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Self {
        Self::compare(lhs, Comparison::Ne, rhs)
    }
    /// Combines this predicate with `rhs` using logical AND.
    pub fn and(self, rhs: Self) -> Self {
        Self::And(Box::new(self), Box::new(rhs))
    }
    /// Combines this predicate with `rhs` using logical OR.
    pub fn or(self, rhs: Self) -> Self {
        Self::Or(Box::new(self), Box::new(rhs))
    }
    /// Creates a closed-interval predicate.
    pub fn between(value: impl Into<Expr>, lower: impl Into<Expr>, upper: impl Into<Expr>) -> Self {
        Self::between_with(value, lower, upper, IntervalClosure::Closed)
    }
    /// Creates an interval predicate with an explicit endpoint policy.
    pub fn between_with(
        value: impl Into<Expr>,
        lower: impl Into<Expr>,
        upper: impl Into<Expr>,
        closure: IntervalClosure,
    ) -> Self {
        Self::Between {
            value: value.into(),
            lower: lower.into(),
            upper: upper.into(),
            closure,
        }
    }
}

impl std::ops::Not for Predicate {
    type Output = Self;
    fn not(self) -> Self::Output {
        Self::Not(Box::new(self))
    }
}

/// Validated, monotonically increasing bin edges.
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(transparent)]
pub struct BinSpec {
    axis: BinningAxis,
}

impl<'de> Deserialize<'de> for BinSpec {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let edges = Vec::<f64>::deserialize(deserializer)?;
        Self::edges(edges).map_err(serde::de::Error::custom)
    }
}

impl BinSpec {
    /// Creates `count` uniformly spaced bins spanning `[min, max]`.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when `count` is zero or the bounds are
    /// non-finite or not increasing.
    pub fn uniform(count: usize, min: f64, max: f64) -> RuntimeResult<Self> {
        Ok(Self {
            axis: BinningAxis::uniform(count, min, max)
                .map_err(|error| query_error(error.to_string()))?,
        })
    }

    /// Creates bins from explicit, strictly increasing finite edges.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when fewer than two edges are supplied or an
    /// edge is non-finite or not strictly increasing.
    pub fn edges(edges: impl IntoIterator<Item = f64>) -> RuntimeResult<Self> {
        Ok(Self {
            axis: BinningAxis::new(edges).map_err(|error| query_error(error.to_string()))?,
        })
    }

    /// Returns the number of bins.
    pub fn bin_count(&self) -> usize {
        self.axis.bin_count()
    }
    /// Returns the validated bin edges.
    pub fn edges_slice(&self) -> &[f64] {
        self.axis.edges()
    }

    fn index(&self, value: f64) -> Option<usize> {
        self.axis.index(value, FinalUpperEdge::Inclusive)
    }
}

/// A lazily filtered dataset corresponding to one bin.
#[derive(Clone)]
pub struct DatasetBin {
    index: usize,
    lower: f64,
    upper: f64,
    dataset: Dataset,
}

impl DatasetBin {
    /// Returns the zero-based bin index.
    pub fn index(&self) -> usize {
        self.index
    }
    /// Returns the bin's lower edge.
    pub fn lower(&self) -> f64 {
        self.lower
    }
    /// Returns the bin's upper edge.
    pub fn upper(&self) -> f64 {
        self.upper
    }
    /// Returns the dataset containing events in this bin.
    pub fn dataset(&self) -> &Dataset {
        &self.dataset
    }
    /// Consumes the bin and returns its dataset.
    pub fn into_dataset(self) -> Dataset {
        self.dataset
    }
}

/// Expression-based query operations for datasets.
pub trait DatasetExprExt {
    /// Validates real scalar query expressions without reading dataset events.
    ///
    /// # Errors
    /// Returns an error for empty, parameter-dependent, or non-real expressions,
    /// or preparation failure.
    fn validate_real_expressions(
        &self,
        expressions: &[Expr],
        execution: &Execution,
    ) -> RuntimeResult<()>;
    /// Evaluates a scalar expression for every event.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when compilation, dataset reading, or
    /// evaluation fails, or the expression is not scalar.
    fn evaluate_expr(&self, expr: &Expr, execution: &Execution) -> RuntimeResult<Vec<Complex64>>;
    /// Evaluates a real scalar expression for every event.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when compilation, dataset reading, or
    /// evaluation fails, or the expression is not real scalar-valued.
    fn evaluate_real(&self, expr: &Expr, execution: &Execution) -> RuntimeResult<Vec<f64>>;
    /// Visits real scalar expression values in bounded event chunks.
    /// The offset is the first event's global row number and expressions retain
    /// their requested order within each callback.
    ///
    /// # Errors
    /// Returns an error for zero chunk size, invalid expressions, dataset access,
    /// evaluation, or a callback failure.
    fn visit_real_chunks(
        &self,
        expressions: &[Expr],
        execution: &Execution,
        chunk_size: usize,
        consume: impl FnMut(usize, &[Vec<f64>]) -> RuntimeResult<()>,
    ) -> RuntimeResult<()>;
    /// Evaluates and fills one weighted histogram in a bounded dataset traversal.
    ///
    /// Dataset event weights are used when `event_weights` is true. When
    /// `weight` is supplied, its real scalar value multiplies the selected
    /// event weight (or unit weight when event weights are disabled).
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when bin validation, expression preparation,
    /// dataset reading, evaluation, or histogram filling fails.
    fn histogram(
        &self,
        expr: &Expr,
        bins: BinSpec,
        event_weights: bool,
        weight: Option<&Expr>,
        execution: &Execution,
    ) -> RuntimeResult<Histogram>;
    /// Evaluates ordered scalar axes and fills one bounded row-major joint histogram.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] for missing or mismatched axes, invalid edges,
    /// non-scalar expressions, evaluation failures, or source read failures.
    fn joint_histogram(
        &self,
        axes: &[Expr],
        bins: Vec<BinSpec>,
        event_weights: bool,
        weight: Option<&Expr>,
        execution: &Execution,
    ) -> RuntimeResult<JointHistogram>;
    /// Creates a lazily filtered dataset containing events that satisfy `predicate`.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when predicate compilation or evaluation
    /// fails, or its expression is not real scalar-valued.
    fn select(&self, predicate: &Predicate, execution: &Execution) -> RuntimeResult<Dataset>;
    /// Partitions the dataset in one pass according to an expression and bin specification.
    ///
    /// # Errors
    ///
    /// Returns [`RuntimeError`] when expression compilation or evaluation
    /// fails, or the expression is not real scalar-valued.
    fn bin_by(
        &self,
        expr: &Expr,
        bins: BinSpec,
        execution: &Execution,
    ) -> RuntimeResult<Vec<DatasetBin>>;
}

impl DatasetExprExt for Dataset {
    fn validate_real_expressions(
        &self,
        expressions: &[Expr],
        execution: &Execution,
    ) -> RuntimeResult<()> {
        if expressions.is_empty() {
            return Err(query_error(
                "real expression validation needs at least one expression",
            ));
        }
        QueryExprSet::prepare(expressions.to_vec(), execution, true)?;
        Ok(())
    }

    fn evaluate_expr(&self, expr: &Expr, execution: &Execution) -> RuntimeResult<Vec<Complex64>> {
        let query = QueryExprSet::prepare(vec![expr.clone()], execution, false)?;
        let mut output = Vec::new();
        for batch in self.batches().map_err(data_error)? {
            output.extend(
                query.evaluate_batch(&batch.map_err(data_error)?)?[0]
                    .iter()
                    .copied(),
            );
        }
        Ok(output)
    }

    fn evaluate_real(&self, expr: &Expr, execution: &Execution) -> RuntimeResult<Vec<f64>> {
        let query = QueryExprSet::prepare(vec![expr.clone()], execution, true)?;
        let mut output = Vec::new();
        for batch in self.batches().map_err(data_error)? {
            output.extend(
                query.evaluate_batch(&batch.map_err(data_error)?)?[0]
                    .iter()
                    .copied()
                    .map(|v| v.re),
            );
        }
        Ok(output)
    }

    fn visit_real_chunks(
        &self,
        expressions: &[Expr],
        execution: &Execution,
        chunk_size: usize,
        mut consume: impl FnMut(usize, &[Vec<f64>]) -> RuntimeResult<()>,
    ) -> RuntimeResult<()> {
        if chunk_size == 0 || expressions.is_empty() {
            return Err(query_error(
                "real chunk evaluation needs expressions and positive chunk size",
            ));
        }
        let query = QueryExprSet::prepare(expressions.to_vec(), execution, true)?;
        let mut offset = 0;
        for batch in self.batches().map_err(data_error)? {
            let batch = batch.map_err(data_error)?;
            for start in (0..batch.len()).step_by(chunk_size) {
                let end = (start + chunk_size).min(batch.len());
                let values = query
                    .evaluate_batch(&batch.slice(start, end))?
                    .into_iter()
                    .map(|column| column.into_iter().map(|value| value.re).collect::<Vec<_>>())
                    .collect::<Vec<_>>();
                consume(offset, &values)?;
                offset += end - start;
            }
        }
        Ok(())
    }

    fn histogram(
        &self,
        expr: &Expr,
        bins: BinSpec,
        event_weights: bool,
        weight: Option<&Expr>,
        execution: &Execution,
    ) -> RuntimeResult<Histogram> {
        let mut expressions = vec![expr.clone()];
        expressions.extend(weight.cloned());
        let query = QueryExprSet::prepare(expressions, execution, true)?;
        let mut histogram = Histogram::empty_with_edges(bins.edges_slice().to_vec())
            .map_err(|error| query_error(error.to_string()))?;

        for batch in self.batches().map_err(data_error)? {
            let batch = batch.map_err(data_error)?;
            let values = query.evaluate_batch(&batch)?;
            for row in 0..batch.len() {
                let base_weight = if event_weights {
                    batch.weights_at(row)
                } else {
                    1.0
                };
                let custom_weight = values.get(1).map_or(1.0, |weights| weights[row].re);
                histogram
                    .fill_weighted(values[0][row].re, base_weight * custom_weight)
                    .map_err(|error| query_error(error.to_string()))?;
            }
        }
        Ok(histogram)
    }

    fn joint_histogram(
        &self,
        axes: &[Expr],
        bins: Vec<BinSpec>,
        event_weights: bool,
        weight: Option<&Expr>,
        execution: &Execution,
    ) -> RuntimeResult<JointHistogram> {
        if axes.is_empty() || axes.len() != bins.len() {
            return Err(query_error(
                "joint histogram requires one bin specification per non-empty ordered axis",
            ));
        }
        let edge_vectors = bins
            .iter()
            .map(|bins| bins.edges_slice().to_vec())
            .collect();
        let mut histogram =
            JointHistogram::empty(edge_vectors).map_err(|error| query_error(error.to_string()))?;
        let mut expressions = axes.to_vec();
        expressions.extend(weight.cloned());
        let query = QueryExprSet::prepare(expressions, execution, true)?;
        let mut coordinates = vec![0.0; axes.len()];
        for batch in self.batches().map_err(data_error)? {
            let batch = batch.map_err(data_error)?;
            let values = query.evaluate_batch(&batch)?;
            for row in 0..batch.len() {
                for (axis, coordinate) in coordinates.iter_mut().enumerate() {
                    *coordinate = values[axis][row].re;
                }
                let base_weight = if event_weights {
                    batch.weights_at(row)
                } else {
                    1.0
                };
                let custom_weight = values
                    .get(axes.len())
                    .map_or(1.0, |weights| weights[row].re);
                histogram
                    .fill_weighted(&coordinates, base_weight * custom_weight)
                    .map_err(|error| query_error(error.to_string()))?;
            }
        }
        Ok(histogram)
    }

    fn select(&self, predicate: &Predicate, execution: &Execution) -> RuntimeResult<Dataset> {
        let compiled = CompiledPredicate::prepare(predicate, execution)?;
        Ok(self.with_derived_source(QuerySource {
            source: self.clone(),
            filter: QueryFilter::Predicate(Arc::new(compiled)),
        }))
    }

    fn bin_by(
        &self,
        expr: &Expr,
        bins: BinSpec,
        execution: &Execution,
    ) -> RuntimeResult<Vec<DatasetBin>> {
        let query = QueryExprSet::prepare(vec![expr.clone()], execution, true)?;
        let schema = self.schema().map_err(data_error)?;
        let mut partitions = vec![Vec::new(); bins.bin_count()];
        for batch in self.batches().map_err(data_error)? {
            let batch = batch.map_err(data_error)?;
            let mut rows = vec![Vec::new(); bins.bin_count()];
            for (row, value) in query.evaluate_batch(&batch)?[0].iter().copied().enumerate() {
                if let Some(index) = bins.index(value.re) {
                    rows[index].push(row);
                }
            }
            for (partition, rows) in partitions.iter_mut().zip(rows) {
                if !rows.is_empty() {
                    partition.push(batch.select(&rows));
                }
            }
        }

        partitions
            .into_iter()
            .enumerate()
            .map(|(index, batches)| {
                let source = if batches.is_empty() {
                    MemorySource::empty(Arc::clone(&schema))
                } else {
                    MemorySource::from_batches(batches).map_err(data_error)?
                };
                Ok(DatasetBin {
                    index,
                    lower: bins.edges_slice()[index],
                    upper: bins.edges_slice()[index + 1],
                    dataset: self.with_derived_source(source),
                })
            })
            .collect()
    }
}

struct QueryExpr {
    model: PreparedModel,
    params: laddu_expr::parameters::ParamValues,
    outputs: Vec<laddu_expr::ExprId>,
}

struct QueryExprSet {
    shared: QueryExprStorage,
    outputs: usize,
}

enum QueryExprStorage {
    Shared(QueryExpr),
    Separate(Vec<QueryExpr>),
}

impl QueryExprSet {
    fn prepare(
        expressions: Vec<Expr>,
        execution: &Execution,
        require_real: bool,
    ) -> RuntimeResult<Self> {
        let expression_count = expressions.len();
        let compiled = CompiledQuery::from_exprs(expressions.clone())
            .map_err(|error| query_error(error.to_string()))?;
        let model = compiled.model();
        let outputs = compiled.outputs();
        if outputs.len() != expression_count {
            return Err(query_error(
                "compiled query output count changed during lowering",
            ));
        }
        for element in outputs {
            let value_kind = model
                .node_facts(*element)
                .map(|facts| facts.value_kind)
                .ok_or_else(|| query_error("compiled expression facts are incomplete"))?;
            if value_kind != ValueKind::Real
                && (require_real || !matches!(value_kind, ValueKind::Complex))
            {
                return Err(query_error(if require_real {
                    "this dataset operation requires a real-valued expression"
                } else {
                    "dataset expressions must be scalar"
                }));
            }
        }
        if model.params().n_free() != 0 {
            return Err(query_error(
                "dataset expressions cannot contain free parameters",
            ));
        }
        let params = model.params().default_values();
        let plan = match PreparedModel::prepare(model, execution) {
            Ok(plan) => QueryExprStorage::Shared(QueryExpr {
                model: plan,
                params,
                outputs: outputs.to_vec(),
            }),
            Err(shared_error) => {
                if !may_fallback_to_scalar(execution, &shared_error) {
                    return Err(shared_error);
                }
                let separate = expressions
                    .iter()
                    .map(|expr| QueryExpr::prepare(expr, execution, require_real))
                    .collect::<RuntimeResult<Vec<_>>>();
                QueryExprStorage::Separate(separate?)
            }
        };
        Ok(Self {
            shared: plan,
            outputs: expression_count,
        })
    }

    fn evaluate_batch(&self, batch: &EventBatch) -> RuntimeResult<Vec<Vec<Complex64>>> {
        let values = match &self.shared {
            QueryExprStorage::Shared(query) => {
                query
                    .model
                    .evaluate_batch_outputs(&query.params, batch, &query.outputs)?
            }
            QueryExprStorage::Separate(queries) => queries
                .iter()
                .map(|query| query.evaluate_batch(batch))
                .collect::<RuntimeResult<Vec<_>>>()?,
        };
        if values.len() != self.outputs {
            return Err(query_error(
                "compiled query returned an unexpected output count",
            ));
        }
        Ok(values)
    }
}

impl QueryExpr {
    fn prepare(expr: &Expr, execution: &Execution, require_real: bool) -> RuntimeResult<Self> {
        if expr.shape().map_err(|e| query_error(e.to_string()))? != laddu_expr::ExprShape::Scalar {
            return Err(query_error("dataset expressions must be scalar"));
        }
        let compiled = laddu_compile::CompiledModel::from_expr(expr)
            .map_err(|error| query_error(error.to_string()))?;
        let value_kind = compiled
            .node_facts(compiled.graph().root())
            .map(|facts| facts.value_kind)
            .ok_or_else(|| query_error("compiled expression facts are incomplete"))?;
        if require_real && value_kind == ValueKind::Complex {
            return Err(query_error(
                "this dataset operation requires a real-valued expression",
            ));
        }
        if compiled.params().n_free() != 0 {
            return Err(query_error(
                "dataset expressions cannot contain free parameters",
            ));
        }
        let params = compiled.params().default_values();
        let model = PreparedModel::prepare(&compiled, execution)?;
        Ok(Self {
            model,
            params,
            outputs: Vec::new(),
        })
    }

    fn evaluate_batch(&self, batch: &EventBatch) -> RuntimeResult<Vec<Complex64>> {
        self.model.evaluate_batch(&self.params, batch)
    }
}

struct CompiledPredicate {
    expressions: QueryExprSet,
    program: PredicateProgram,
}

enum PredicateProgram {
    Compare {
        lhs: usize,
        op: Comparison,
        rhs: usize,
    },
    And(Box<Self>, Box<Self>),
    Or(Box<Self>, Box<Self>),
    Not(Box<Self>),
    Between {
        value: usize,
        lower: usize,
        upper: usize,
        closure: IntervalClosure,
    },
}

impl CompiledPredicate {
    fn prepare(predicate: &Predicate, execution: &Execution) -> RuntimeResult<Self> {
        let mut expressions = Vec::new();
        let program = Self::compile_program(predicate, &mut expressions);
        Ok(Self {
            expressions: QueryExprSet::prepare(expressions, execution, true)?,
            program,
        })
    }

    fn compile_program(predicate: &Predicate, expressions: &mut Vec<Expr>) -> PredicateProgram {
        let leaf = |expr: &Expr, expressions: &mut Vec<Expr>| {
            let index = expressions.len();
            expressions.push(expr.clone());
            index
        };
        match predicate {
            Predicate::Compare { lhs, op, rhs } => PredicateProgram::Compare {
                lhs: leaf(lhs, expressions),
                op: *op,
                rhs: leaf(rhs, expressions),
            },
            Predicate::And(lhs, rhs) => PredicateProgram::And(
                Box::new(Self::compile_program(lhs, expressions)),
                Box::new(Self::compile_program(rhs, expressions)),
            ),
            Predicate::Or(lhs, rhs) => PredicateProgram::Or(
                Box::new(Self::compile_program(lhs, expressions)),
                Box::new(Self::compile_program(rhs, expressions)),
            ),
            Predicate::Not(inner) => {
                PredicateProgram::Not(Box::new(Self::compile_program(inner, expressions)))
            }
            Predicate::Between {
                value,
                lower,
                upper,
                closure,
            } => PredicateProgram::Between {
                value: leaf(value, expressions),
                lower: leaf(lower, expressions),
                upper: leaf(upper, expressions),
                closure: *closure,
            },
        }
    }

    fn evaluate_batch(&self, batch: &EventBatch) -> RuntimeResult<Vec<usize>> {
        let values = self.expressions.evaluate_batch(batch)?;
        Ok((0..batch.len())
            .filter(|row| Self::evaluate_row(&self.program, &values, *row))
            .collect())
    }

    fn evaluate_row(program: &PredicateProgram, values: &[Vec<Complex64>], row: usize) -> bool {
        match program {
            PredicateProgram::Compare { lhs, op, rhs } => {
                compare(values[*lhs][row].re, *op, values[*rhs][row].re)
            }
            PredicateProgram::And(lhs, rhs) => {
                Self::evaluate_row(lhs, values, row) && Self::evaluate_row(rhs, values, row)
            }
            PredicateProgram::Or(lhs, rhs) => {
                Self::evaluate_row(lhs, values, row) || Self::evaluate_row(rhs, values, row)
            }
            PredicateProgram::Not(inner) => !Self::evaluate_row(inner, values, row),
            PredicateProgram::Between {
                value,
                lower,
                upper,
                closure,
            } => {
                let lower_op = match closure {
                    IntervalClosure::Open | IntervalClosure::RightClosed => Comparison::Gt,
                    IntervalClosure::LeftClosed | IntervalClosure::Closed => Comparison::Ge,
                };
                let upper_op = match closure {
                    IntervalClosure::Open | IntervalClosure::LeftClosed => Comparison::Lt,
                    IntervalClosure::RightClosed | IntervalClosure::Closed => Comparison::Le,
                };
                compare(values[*value][row].re, lower_op, values[*lower][row].re)
                    && compare(values[*value][row].re, upper_op, values[*upper][row].re)
            }
        }
    }
}

fn compare(lhs: f64, op: Comparison, rhs: f64) -> bool {
    if lhs.is_nan() || rhs.is_nan() {
        return false;
    }
    match op {
        Comparison::Lt => lhs < rhs,
        Comparison::Le => lhs <= rhs,
        Comparison::Gt => lhs > rhs,
        Comparison::Ge => lhs >= rhs,
        Comparison::Eq => lhs == rhs,
        Comparison::Ne => lhs != rhs,
    }
}

#[derive(Clone)]
struct QuerySource {
    source: Dataset,
    filter: QueryFilter,
}

#[derive(Clone)]
enum QueryFilter {
    Predicate(Arc<CompiledPredicate>),
}

impl EventSource for QuerySource {
    fn schema(&self) -> LadduDataResult<Arc<Schema>> {
        self.source.schema()
    }

    fn capabilities(&self) -> SourceCapabilities {
        let source = self.source.capabilities();
        SourceCapabilities {
            exact_len: false,
            exact_weighted_total: false,
            random_access: false,
            deterministic_partitioning: source.deterministic_partitioning,
            predicate_pushdown: false,
            projection_pushdown: false,
            streaming: true,
        }
    }

    fn batches(&self, plan: ReadPlan) -> LadduDataResult<EventBatchIter> {
        let batches = self.source.stream_with_plan(plan)?;
        let filter = self.filter.clone();
        Ok(Box::new(batches.filter_map(move |batch| {
            let batch = match batch {
                Ok(batch) => batch,
                Err(error) => return Some(Err(error)),
            };
            let rows = match filter.rows(&batch) {
                Ok(rows) => rows,
                Err(error) => return Some(Err(LadduDataError::Source(error.to_string()))),
            };
            (!rows.is_empty()).then(|| Ok(batch.select(&rows)))
        })))
    }
}

impl QueryFilter {
    fn rows(&self, batch: &EventBatch) -> RuntimeResult<Vec<usize>> {
        match self {
            Self::Predicate(predicate) => predicate.evaluate_batch(batch),
        }
    }
}

fn query_error(message: impl Into<String>) -> RuntimeError {
    RuntimeError::InvalidShape {
        index: 0,
        message: message.into(),
    }
}

fn may_fallback_to_scalar(execution: &Execution, error: &RuntimeError) -> bool {
    let cpu_f32 = matches!(
        error,
        RuntimeError::Execution(crate::ExecutionError::UnsupportedCpuF32Model)
    );
    #[cfg(feature = "wgpu")]
    {
        cpu_f32 || (execution.wgpu_context().is_some() && matches!(error, RuntimeError::Wgpu(_)))
    }
    #[cfg(not(feature = "wgpu"))]
    {
        let _ = execution;
        cpu_f32
    }
}

fn data_error(error: impl ToString) -> RuntimeError {
    RuntimeError::Data(error.to_string())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{CpuOptions, Device, ExecutionOptions, Precision};
    #[cfg(feature = "jit")]
    use crate::{JitPolicy, ThreadPolicy};
    use laddu_compile::CompiledModel;
    use laddu_data::{
        data::{EventBatch, OwnedEvent},
        io::{EventSource, ReadPlan, SourceCapabilities, memory::MemorySource},
        schema::Schema,
    };
    use laddu_expr::{complex, event_scalar};
    use std::sync::atomic::{AtomicUsize, Ordering};

    #[test]
    fn bin_spec_roundtrip_preserves_validation() {
        let bins = BinSpec::edges([-1.0, 0.0, 2.0]).unwrap();
        let json = serde_json::to_string(&bins).unwrap();
        assert_eq!(serde_json::from_str::<BinSpec>(&json).unwrap(), bins);
        assert!(serde_json::from_str::<BinSpec>("[0.0,0.0]").is_err());
    }

    #[derive(Clone)]
    struct CountingSource {
        inner: MemorySource,
        reads: Arc<AtomicUsize>,
    }

    impl EventSource for CountingSource {
        fn schema(&self) -> LadduDataResult<Arc<Schema>> {
            EventSource::schema(&self.inner)
        }

        fn capabilities(&self) -> SourceCapabilities {
            self.inner.capabilities()
        }

        fn batches(&self, plan: ReadPlan) -> LadduDataResult<EventBatchIter> {
            self.reads.fetch_add(1, Ordering::Relaxed);
            self.inner.batches(plan)
        }
    }

    #[derive(Clone)]
    struct FailingSource {
        schema: Arc<Schema>,
    }

    impl EventSource for FailingSource {
        fn schema(&self) -> LadduDataResult<Arc<Schema>> {
            Ok(Arc::clone(&self.schema))
        }

        fn capabilities(&self) -> SourceCapabilities {
            SourceCapabilities {
                exact_len: false,
                exact_weighted_total: false,
                random_access: false,
                deterministic_partitioning: true,
                predicate_pushdown: false,
                projection_pushdown: false,
                streaming: true,
            }
        }

        fn batches(&self, _plan: ReadPlan) -> LadduDataResult<EventBatchIter> {
            Ok(Box::new(std::iter::once(Err(LadduDataError::Source(
                "query source failed".into(),
            )))))
        }
    }

    fn capability_tuple(
        capabilities: SourceCapabilities,
    ) -> (bool, bool, bool, bool, bool, bool, bool) {
        (
            capabilities.exact_len,
            capabilities.exact_weighted_total,
            capabilities.random_access,
            capabilities.deterministic_partitioning,
            capabilities.predicate_pushdown,
            capabilities.projection_pushdown,
            capabilities.streaming,
        )
    }

    fn dataset() -> Dataset {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        Dataset::from_events(
            schema,
            [
                OwnedEvent::weighted(vec![], vec![-1.0], 0.5),
                OwnedEvent::weighted(vec![], vec![0.0], 1.0),
                OwnedEvent::weighted(vec![], vec![1.0], 1.5),
                OwnedEvent::weighted(vec![], vec![2.0], 2.0),
            ],
        )
        .unwrap()
    }

    #[test]
    fn evaluates_selects_and_bins_dataset_expressions() {
        let dataset = dataset().chunked(1).unwrap();
        let execution = Execution::default();
        let x = event_scalar("x");
        assert_eq!(
            dataset.evaluate_real(&x, &execution).unwrap(),
            vec![-1.0, 0.0, 1.0, 2.0]
        );

        let selected = dataset
            .select(
                &Predicate::ge(x.clone(), 0.0).and(Predicate::lt(x.clone(), 2.0)),
                &execution,
            )
            .unwrap();
        assert_eq!(
            selected.map_events(|event| event.scalar(0)).unwrap(),
            vec![0.0, 1.0]
        );
        assert_eq!(selected.sum_weights().unwrap(), 2.5);

        let bins = dataset
            .bin_by(&x, BinSpec::uniform(2, 0.0, 2.0).unwrap(), &execution)
            .unwrap();
        assert_eq!(bins.len(), 2);
        assert_eq!(
            bins[0]
                .dataset()
                .map_events(|event| event.scalar(0))
                .unwrap(),
            vec![0.0]
        );
        assert_eq!(
            bins[1]
                .dataset()
                .map_events(|event| event.scalar(0))
                .unwrap(),
            vec![1.0, 2.0]
        );
    }

    #[test]
    fn dataset_histogram_uses_event_weights_and_excludes_the_final_upper_edge() {
        let histogram = dataset()
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0, 2.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(histogram.counts(), [1.0, 1.5]);
        assert_eq!(histogram.sum_squared_weights(), [1.0, 2.25]);
        assert_eq!(histogram.underflow(), 0.5);
        assert_eq!(histogram.underflow_sum_squared_weights(), Some(0.25));
        assert_eq!(histogram.overflow(), 2.0);
        assert_eq!(histogram.overflow_sum_squared_weights(), Some(4.0));
    }

    #[test]
    fn dataset_joint_histogram_uses_row_major_bins_and_aggregates_invalid_events() {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x", "y"], true).unwrap());
        let dataset = Dataset::from_events(
            schema,
            [
                OwnedEvent::weighted(vec![], vec![0.0, 10.0], 1.0),
                OwnedEvent::weighted(vec![], vec![1.0, 10.0], -2.0),
                OwnedEvent::weighted(vec![], vec![0.0, 20.0], 3.0),
                OwnedEvent::weighted(vec![], vec![f64::NAN, 10.0], 4.0),
            ],
        )
        .unwrap();

        let histogram = dataset
            .joint_histogram(
                &[event_scalar("x"), event_scalar("y")],
                vec![
                    BinSpec::edges([0.0, 1.0, 2.0]).unwrap(),
                    BinSpec::edges([0.0, 15.0, 25.0]).unwrap(),
                ],
                true,
                None,
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(histogram.shape(), [2, 2]);
        assert_eq!(histogram.values(), [1.0, 3.0, -2.0, 0.0]);
        assert_eq!(histogram.sum_squared_weights(), [1.0, 9.0, 4.0, 0.0]);
        assert_eq!(histogram.diagnostics().nonfinite_count(), 1);
        assert_eq!(histogram.diagnostics().out_of_range_count(), 0);
    }

    #[test]
    fn dataset_histogram_multiplies_the_selected_base_and_custom_weights() {
        let x = event_scalar("x");
        let bins = BinSpec::edges([-2.0, 0.0, 2.0]).unwrap();
        let execution = Execution::default();

        let weighted = dataset()
            .histogram(&x, bins.clone(), true, Some(&(x.clone() + 2.0)), &execution)
            .unwrap();
        assert_eq!(weighted.counts(), [0.5, 6.5]);
        assert_eq!(weighted.sum_squared_weights(), [0.25, 24.25]);
        assert_eq!(weighted.overflow(), 8.0);

        let custom_only = dataset()
            .histogram(&x, bins, false, Some(&(x.clone() + 2.0)), &execution)
            .unwrap();
        assert_eq!(custom_only.counts(), [1.0, 5.0]);
        assert_eq!(custom_only.sum_squared_weights(), [1.0, 13.0]);
        assert_eq!(custom_only.overflow(), 4.0);
    }

    #[test]
    fn dataset_histogram_preserves_view_source_and_memory_semantics() {
        let source = dataset();
        let batch = source.batches().unwrap().next().unwrap().unwrap();
        let reads = Arc::new(AtomicUsize::new(0));
        let counted = Dataset::new(CountingSource {
            inner: MemorySource::new(batch),
            reads: Arc::clone(&reads),
        })
        .streaming()
        .chunked(1)
        .unwrap();
        let selected = counted
            .select(
                &Predicate::ge(event_scalar("x"), 0.0),
                &Execution::default(),
            )
            .unwrap();

        let histogram = selected
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0, 2.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(histogram.counts(), [1.0, 1.5]);
        assert_eq!(histogram.overflow(), 2.0);
        assert_eq!(reads.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn dataset_histogram_matches_across_memory_policies_and_chunking() {
        let source = dataset();
        let bins = BinSpec::edges([0.0, 1.0, 2.0]).unwrap();
        let execution = Execution::default();
        let expected = source
            .clone()
            .resident()
            .histogram(&event_scalar("x"), bins.clone(), true, None, &execution)
            .unwrap();

        for candidate in [
            source.clone().streaming(),
            source.clone().resident().chunked(1).unwrap(),
            source.streaming().chunked(2).unwrap(),
        ] {
            let actual = candidate
                .histogram(&event_scalar("x"), bins.clone(), true, None, &execution)
                .unwrap();
            assert_eq!(actual, expected);
        }
    }

    #[cfg(feature = "jit")]
    #[test]
    fn dataset_histogram_matches_cpu_interpreter_and_jit_backends() {
        let source = dataset();
        let bins = BinSpec::edges([0.0, 1.0, 2.0]).unwrap();
        let execution = |jit| {
            Execution::local(ExecutionOptions {
                device: Device::Cpu(CpuOptions {
                    threads: ThreadPolicy::Serial,
                    jit,
                }),
                precision: Precision::F64,
                ..ExecutionOptions::default()
            })
            .unwrap()
        };
        let interpreted = source
            .histogram(
                &event_scalar("x"),
                bins.clone(),
                true,
                None,
                &execution(JitPolicy::Disabled),
            )
            .unwrap();
        let compiled = source
            .histogram(
                &event_scalar("x"),
                bins,
                true,
                None,
                &execution(JitPolicy::Enabled),
            )
            .unwrap();

        assert_eq!(compiled, interpreted);
    }

    #[test]
    fn dataset_histogram_reports_invalid_values_and_source_failures() {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        let nonfinite = Dataset::from_events(
            Arc::clone(&schema),
            [OwnedEvent::weighted(vec![], vec![f64::NAN], 1.0)],
        )
        .unwrap();
        let error = nonfinite
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap_err();
        assert!(
            error.to_string().contains("expected finite, got NaN"),
            "unexpected error: {error}",
        );

        let source_error = Dataset::new(FailingSource { schema })
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap_err();
        assert!(source_error.to_string().contains("query source failed"));
    }

    #[test]
    fn empty_dataset_histogram_is_a_valid_empirical_histogram() {
        let histogram = dataset()
            .empty_derived()
            .unwrap()
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0, 3.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(histogram.counts(), [0.0, 0.0]);
        assert_eq!(histogram.sum_squared_weights(), [0.0, 0.0]);
        assert_eq!(histogram.underflow_sum_squared_weights(), Some(0.0));
        assert_eq!(histogram.overflow_sum_squared_weights(), Some(0.0));
    }

    #[test]
    fn cancelling_dataset_weights_keep_their_squared_weight_uncertainty() {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        let cancelling = Dataset::from_events(
            schema,
            [
                OwnedEvent::weighted(vec![], vec![0.5], 1.0),
                OwnedEvent::weighted(vec![], vec![0.5], -1.0),
            ],
        )
        .unwrap();
        let histogram = cancelling
            .histogram(
                &event_scalar("x"),
                BinSpec::edges([0.0, 1.0]).unwrap(),
                true,
                None,
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(histogram.counts(), [0.0]);
        assert_eq!(histogram.sum_squared_weights(), [2.0]);
        assert_eq!(histogram.errors(), [2.0_f64.sqrt()]);
    }

    #[test]
    fn empty_batches_are_valid_query_inputs() {
        let execution = Execution::default();
        let x = event_scalar("x");

        let empty_batch_schema =
            Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        let empty_batch = Dataset::from_batch(
            EventBatch::from_events(empty_batch_schema, std::iter::empty::<OwnedEvent>()).unwrap(),
        );
        for empty in [empty_batch, dataset().empty_derived().unwrap()] {
            assert!(empty.evaluate_real(&x, &execution).unwrap().is_empty());
            assert!(
                empty
                    .select(&Predicate::ge(x.clone(), 0.0), &execution)
                    .unwrap()
                    .map_events(|event| event.scalar(0))
                    .unwrap()
                    .is_empty()
            );
        }
    }

    #[test]
    fn event_column_nan_comparisons_are_false() {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        let dataset = Dataset::from_events(
            schema,
            [
                OwnedEvent::weighted(vec![], vec![f64::NAN], 1.0),
                OwnedEvent::weighted(vec![], vec![0.0], 1.0),
                OwnedEvent::weighted(vec![], vec![1.0], 1.0),
            ],
        )
        .unwrap();
        let x = event_scalar("x");
        let selected = dataset
            .select(&Predicate::ne(x, 0.0), &Execution::default())
            .unwrap();

        assert_eq!(
            selected.map_events(|event| event.scalar(0)).unwrap(),
            vec![1.0]
        );
    }

    #[test]
    fn query_propagates_source_batch_errors() {
        let schema = Arc::new(Schema::new(std::iter::empty::<&str>(), ["x"], true).unwrap());
        let dataset = Dataset::new(FailingSource { schema });

        let error = dataset
            .evaluate_real(&event_scalar("x"), &Execution::default())
            .unwrap_err();
        assert!(
            matches!(error, RuntimeError::Data(message) if message.contains("query source failed"))
        );
    }

    #[test]
    fn all_empty_bins_retain_valid_empty_derived_sources() {
        let source = dataset();
        let before = capability_tuple(source.capabilities());
        let bins = source
            .bin_by(
                &event_scalar("x"),
                BinSpec::edges([10.0, 20.0, 30.0]).unwrap(),
                &Execution::default(),
            )
            .unwrap();

        assert_eq!(capability_tuple(source.capabilities()), before);
        assert_eq!(bins.len(), 2);
        for bin in bins {
            assert_eq!(bin.dataset().num_events().unwrap(), Some(0));
            assert!(
                bin.dataset()
                    .evaluate_real(&event_scalar("x"), &Execution::default())
                    .unwrap()
                    .is_empty()
            );
        }
    }

    #[test]
    fn traversing_all_bins_reads_the_source_once() {
        let reads = Arc::new(AtomicUsize::new(0));
        let source = CountingSource {
            inner: match dataset().batches().unwrap().next().unwrap() {
                Ok(batch) => MemorySource::new(batch),
                Err(error) => panic!("unexpected source error: {error}"),
            },
            reads: Arc::clone(&reads),
        };
        let dataset = Dataset::new(source).chunked(1).unwrap();
        let bins = dataset
            .bin_by(
                &event_scalar("x"),
                BinSpec::uniform(4, -1.0, 3.0).unwrap(),
                &Execution::default(),
            )
            .unwrap();

        let values = bins
            .into_iter()
            .map(|bin| {
                bin.into_dataset()
                    .map_events(|event| event.scalar(0))
                    .unwrap()
            })
            .collect::<Vec<_>>();
        assert_eq!(values, [vec![-1.0], vec![0.0], vec![1.0], vec![2.0]]);
        assert_eq!(reads.load(Ordering::Relaxed), 1);
    }

    #[test]
    fn between_predicates_have_explicit_endpoint_semantics() {
        let dataset = dataset();
        let execution = Execution::default();
        let x = event_scalar("x");

        let closed = dataset
            .select(&Predicate::between(x.clone(), 0.0, 1.0), &execution)
            .unwrap();
        assert_eq!(
            closed.map_events(|event| event.scalar(0)).unwrap(),
            vec![0.0, 1.0]
        );

        let open = dataset
            .select(
                &Predicate::between_with(x, -1.0, 1.0, IntervalClosure::Open),
                &execution,
            )
            .unwrap();
        assert_eq!(open.map_events(|event| event.scalar(0)).unwrap(), vec![0.0]);
    }

    #[test]
    fn real_queries_reject_complex_and_free_parameter_expressions() {
        let dataset = dataset();
        let execution = Execution::default();
        assert!(
            dataset
                .evaluate_real(&complex(1.0, 1.0), &execution)
                .is_err()
        );
        let parameter = Expr::from(laddu_expr::parameters::Parameter::free("p"));
        assert!(dataset.evaluate_expr(&parameter, &execution).is_err());
    }

    #[test]
    fn compiled_query_outputs_preserve_order_and_values() {
        let source = dataset();
        let batch = source.batches().unwrap().next().unwrap().unwrap();
        let x = event_scalar("x");
        let query = QueryExprSet::prepare(
            vec![x.clone() + 1.0, x.clone() * 2.0, x],
            &Execution::default(),
            false,
        )
        .unwrap();
        let values = query.evaluate_batch(&batch).unwrap();
        assert_eq!(
            values[0].iter().map(|v| v.re).collect::<Vec<_>>(),
            [0.0, 1.0, 2.0, 3.0]
        );
        assert_eq!(
            values[1].iter().map(|v| v.re).collect::<Vec<_>>(),
            [-2.0, 0.0, 2.0, 4.0]
        );
        assert_eq!(
            values[2].iter().map(|v| v.re).collect::<Vec<_>>(),
            [-1.0, 0.0, 1.0, 2.0]
        );
    }

    #[test]
    fn repeated_predicate_leaves_are_evaluated_once() {
        let x = event_scalar("x");
        let selected = dataset()
            .select(
                &Predicate::ge(x.clone() + 1.0, 0.0).and(Predicate::lt(x + 1.0, 2.0)),
                &Execution::default(),
            )
            .unwrap();
        assert_eq!(
            selected.map_events(|event| event.scalar(0)).unwrap(),
            [-1.0, 0.0]
        );
    }

    #[test]
    fn f32_queries_match_f64_query_results() {
        let x = event_scalar("x");
        let f64_values = dataset().evaluate_real(&x, &Execution::default()).unwrap();
        let f32_execution = Execution::local(ExecutionOptions {
            device: Device::Cpu(CpuOptions::default()),
            precision: Precision::F32,
            ..ExecutionOptions::default()
        })
        .unwrap();
        let f32_values = dataset().evaluate_real(&x, &f32_execution).unwrap();
        assert_eq!(f32_values, f64_values);

        let bins = BinSpec::edges([-2.0, 0.0, 2.0]).unwrap();
        let f64_histogram = dataset()
            .histogram(&x, bins.clone(), true, None, &Execution::default())
            .unwrap();
        let f32_histogram = dataset()
            .histogram(&x, bins, true, None, &f32_execution)
            .unwrap();
        assert_eq!(f32_histogram, f64_histogram);
    }

    #[test]
    fn bin_edges_validate_and_nan_predicates_are_false() {
        assert!(BinSpec::edges([0.0, 0.0]).is_err());
        assert!(!compare(f64::NAN, Comparison::Ne, 0.0));
    }

    #[test]
    fn selection_is_lazy_and_one_pass_binning_preserves_streaming_policy() {
        let source = dataset();
        let batch = source.batches().unwrap().next().unwrap().unwrap();
        let reads = Arc::new(AtomicUsize::new(0));
        let dataset = Dataset::new(CountingSource {
            inner: MemorySource::new(batch),
            reads: Arc::clone(&reads),
        })
        .streaming();
        let execution = Execution::default();
        let x = event_scalar("x");

        let selected = dataset
            .select(&Predicate::ge(x.clone(), 0.0), &execution)
            .unwrap();
        let bins = dataset
            .bin_by(&x, BinSpec::uniform(2, 0.0, 2.0).unwrap(), &execution)
            .unwrap();
        assert_eq!(reads.load(Ordering::Relaxed), 1);
        assert_eq!(
            selected.cache_storage(),
            laddu_data::data::CacheStorage::Streaming
        );

        assert_eq!(
            selected.map_events(|event| event.scalar(0)).unwrap(),
            vec![0.0, 1.0, 2.0]
        );
        assert_eq!(reads.load(Ordering::Relaxed), 2);
        assert_eq!(
            bins[0]
                .dataset()
                .map_events(|event| event.scalar(0))
                .unwrap(),
            vec![0.0]
        );
        assert_eq!(reads.load(Ordering::Relaxed), 2);
    }

    #[test]
    fn unknown_cardinality_fastest_discovers_and_retains_small_selection() {
        let source = dataset();
        let batch = source.batches().unwrap().next().unwrap().unwrap();
        let reads = Arc::new(AtomicUsize::new(0));
        let dataset = Dataset::new(CountingSource {
            inner: MemorySource::new(batch),
            reads: Arc::clone(&reads),
        });
        let execution = Execution::default();
        let x = event_scalar("x");
        let selected = dataset
            .select(&Predicate::ge(x.clone(), 0.0), &execution)
            .unwrap();
        let compiled = CompiledModel::from_expr(&x).unwrap();
        let params = compiled.params().default_values();
        let model = PreparedModel::prepare(&compiled, &execution).unwrap();
        let prepared = model.prepare_dataset(&execution, &selected).unwrap();

        #[cfg(not(feature = "wgpu"))]
        let crate::PreparedDataset::Cpu(prepared_cpu) = &prepared;
        #[cfg(feature = "wgpu")]
        let crate::PreparedDataset::Cpu(prepared_cpu) = &prepared else {
            panic!("default execution prepares CPU datasets");
        };
        assert_eq!(
            prepared_cpu.stats().storage(),
            laddu_data::data::CacheStorage::Resident
        );
        assert_eq!(prepared_cpu.stats().local_events(), 3);
        assert_eq!(reads.load(Ordering::Relaxed), 2);

        for _ in 0..2 {
            assert_eq!(
                model
                    .reduce(
                        &execution,
                        &params,
                        &prepared,
                        laddu_compile::ReductionPlan::weighted_real(),
                    )
                    .unwrap(),
                5.5
            );
        }
        assert_eq!(reads.load(Ordering::Relaxed), 2);
    }
}