table_sample.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! # TABLESAMPLE Example
//!
//! This example demonstrates implementing SQL `TABLESAMPLE` support using
//! DataFusion's extensibility APIs.
//!
//! This is a working `TABLESAMPLE` implementation that can serve as a starting
//! point for your own projects. It also works as a template for adding other
//! custom SQL operators, covering the full pipeline from parsing to execution.
//!
//! It shows how to:
//!
//! 1. **Parse** TABLESAMPLE syntax via a custom [`RelationPlanner`]
//! 2. **Plan** sampling as a custom logical node ([`TableSamplePlanNode`])
//! 3. **Execute** sampling via a custom physical operator ([`SampleExec`])
//!
//! ## Supported Syntax
//!
//! ```sql
//! -- Bernoulli sampling (each row has N% chance of selection)
//! SELECT * FROM table TABLESAMPLE BERNOULLI(10 PERCENT)
//!
//! -- Fractional sampling (0.0 to 1.0)
//! SELECT * FROM table TABLESAMPLE (0.1)
//!
//! -- Row count limit
//! SELECT * FROM table TABLESAMPLE (100 ROWS)
//!
//! -- Reproducible sampling with a seed
//! SELECT * FROM table TABLESAMPLE (10 PERCENT) REPEATABLE(42)
//! ```
//!
//! ## Architecture
//!
//! ```text
//! ┌─────────────────────────────────────────────────────────────────┐
//! │                         SQL Query                               │
//! │  SELECT * FROM t TABLESAMPLE BERNOULLI(10 PERCENT) REPEATABLE(1)│
//! └─────────────────────────────────────────────────────────────────┘
//!                                │
//!                                ▼
//! ┌─────────────────────────────────────────────────────────────────┐
//! │                    TableSamplePlanner                           │
//! │    (RelationPlanner: parses TABLESAMPLE, creates logical node)  │
//! └─────────────────────────────────────────────────────────────────┘
//!                                │
//!                                ▼
//! ┌─────────────────────────────────────────────────────────────────┐
//! │                   TableSamplePlanNode                           │
//! │         (UserDefinedLogicalNode: stores sampling params)        │
//! └─────────────────────────────────────────────────────────────────┘
//!                                │
//!                                ▼
//! ┌─────────────────────────────────────────────────────────────────┐
//! │                TableSampleExtensionPlanner                      │
//! │       (ExtensionPlanner: creates physical execution plan)       │
//! └─────────────────────────────────────────────────────────────────┘
//!                                │
//!                                ▼
//! ┌─────────────────────────────────────────────────────────────────┐
//! │                        SampleExec                               │
//! │    (ExecutionPlan: performs actual row sampling at runtime)     │
//! └─────────────────────────────────────────────────────────────────┘
//! ```

use std::{
    any::Any,
    fmt::{self, Debug, Formatter},
    hash::{Hash, Hasher},
    pin::Pin,
    sync::Arc,
    task::{Context, Poll},
};

use arrow::datatypes::{Float64Type, Int64Type};
use arrow::{
    array::{ArrayRef, Int32Array, RecordBatch, StringArray, UInt32Array},
    compute,
};
use arrow_schema::SchemaRef;
use futures::{
    ready,
    stream::{Stream, StreamExt},
};
use rand::{Rng, SeedableRng, rngs::StdRng};
use tonic::async_trait;

use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal;
use datafusion::{
    execution::{
        RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder,
        TaskContext, context::QueryPlanner,
    },
    physical_expr::EquivalenceProperties,
    physical_plan::{
        DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
        metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput},
    },
    physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner},
    prelude::*,
};
use datafusion_common::{
    DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err,
    plan_datafusion_err, plan_err,
};
use datafusion_expr::{
    UserDefinedLogicalNode, UserDefinedLogicalNodeCore,
    logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder},
    planner::{
        PlannedRelation, RelationPlanner, RelationPlannerContext, RelationPlanning,
    },
};
use datafusion_sql::sqlparser::ast::{
    self, TableFactor, TableSampleMethod, TableSampleUnit,
};
use insta::assert_snapshot;

// ============================================================================
// Example Entry Point
// ============================================================================

/// Runs the TABLESAMPLE examples demonstrating various sampling techniques.
pub async fn table_sample() -> Result<()> {
    // Build session with custom query planner for physical planning
    let state = SessionStateBuilder::new()
        .with_default_features()
        .with_query_planner(Arc::new(TableSampleQueryPlanner))
        .build();

    let ctx = SessionContext::new_with_state(state);

    // Register custom relation planner for logical planning
    ctx.register_relation_planner(Arc::new(TableSamplePlanner))?;
    register_sample_data(&ctx)?;

    println!("TABLESAMPLE Example");
    println!("===================\n");

    run_examples(&ctx).await
}

async fn run_examples(ctx: &SessionContext) -> Result<()> {
    // Example 1: Baseline - full table scan
    let results = run_example(
        ctx,
        "Example 1: Full table (baseline)",
        "SELECT * FROM sample_data",
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+
    | column1 | column2 |
    +---------+---------+
    | 1       | row_1   |
    | 2       | row_2   |
    | 3       | row_3   |
    | 4       | row_4   |
    | 5       | row_5   |
    | 6       | row_6   |
    | 7       | row_7   |
    | 8       | row_8   |
    | 9       | row_9   |
    | 10      | row_10  |
    +---------+---------+
    ");

    // Example 2: Percentage-based Bernoulli sampling
    // REPEATABLE(seed) ensures deterministic results for snapshot testing
    let results = run_example(
        ctx,
        "Example 2: BERNOULLI percentage sampling",
        "SELECT * FROM sample_data TABLESAMPLE BERNOULLI(30 PERCENT) REPEATABLE(123)",
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+
    | column1 | column2 |
    +---------+---------+
    | 1       | row_1   |
    | 2       | row_2   |
    | 7       | row_7   |
    | 8       | row_8   |
    +---------+---------+
    ");

    // Example 3: Fractional sampling (0.0 to 1.0)
    // REPEATABLE(seed) ensures deterministic results for snapshot testing
    let results = run_example(
        ctx,
        "Example 3: Fractional sampling",
        "SELECT * FROM sample_data TABLESAMPLE (0.5) REPEATABLE(456)",
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+
    | column1 | column2 |
    +---------+---------+
    | 2       | row_2   |
    | 4       | row_4   |
    | 8       | row_8   |
    +---------+---------+
    ");

    // Example 4: Row count limit (deterministic, no seed needed)
    let results = run_example(
        ctx,
        "Example 4: Row count limit",
        "SELECT * FROM sample_data TABLESAMPLE (3 ROWS)",
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+
    | column1 | column2 |
    +---------+---------+
    | 1       | row_1   |
    | 2       | row_2   |
    | 3       | row_3   |
    +---------+---------+
    ");

    // Example 5: Sampling combined with filtering
    let results = run_example(
        ctx,
        "Example 5: Sampling with WHERE clause",
        "SELECT * FROM sample_data TABLESAMPLE (5 ROWS) WHERE column1 > 2",
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+
    | column1 | column2 |
    +---------+---------+
    | 3       | row_3   |
    | 4       | row_4   |
    | 5       | row_5   |
    +---------+---------+
    ");

    // Example 6: Sampling in JOIN queries
    // REPEATABLE(seed) ensures deterministic results for snapshot testing
    let results = run_example(
        ctx,
        "Example 6: Sampling in JOINs",
        r#"SELECT t1.column1, t2.column1, t1.column2, t2.column2
           FROM sample_data t1 TABLESAMPLE (0.7) REPEATABLE(789)
           JOIN sample_data t2 TABLESAMPLE (0.7) REPEATABLE(123)
           ON t1.column1 = t2.column1"#,
    )
    .await?;
    assert_snapshot!(results, @r"
    +---------+---------+---------+---------+
    | column1 | column1 | column2 | column2 |
    +---------+---------+---------+---------+
    | 2       | 2       | row_2   | row_2   |
    | 5       | 5       | row_5   | row_5   |
    | 7       | 7       | row_7   | row_7   |
    | 8       | 8       | row_8   | row_8   |
    | 10      | 10      | row_10  | row_10  |
    +---------+---------+---------+---------+
    ");

    Ok(())
}

/// Helper to run a single example query and capture results.
async fn run_example(ctx: &SessionContext, title: &str, sql: &str) -> Result<String> {
    println!("{title}:\n{sql}\n");
    let df = ctx.sql(sql).await?;
    println!("{}\n", df.logical_plan().display_indent());

    let batches = df.collect().await?;
    let results = arrow::util::pretty::pretty_format_batches(&batches)?.to_string();
    println!("{results}\n");

    Ok(results)
}

/// Register test data: 10 rows with column1=1..10 and column2="row_1".."row_10"
fn register_sample_data(ctx: &SessionContext) -> Result<()> {
    let column1: ArrayRef = Arc::new(Int32Array::from((1..=10).collect::<Vec<i32>>()));
    let column2: ArrayRef = Arc::new(StringArray::from(
        (1..=10).map(|i| format!("row_{i}")).collect::<Vec<_>>(),
    ));
    let batch =
        RecordBatch::try_from_iter(vec![("column1", column1), ("column2", column2)])?;
    ctx.register_batch("sample_data", batch)?;
    Ok(())
}

// ============================================================================
// Logical Planning: TableSamplePlanner + TableSamplePlanNode
// ============================================================================

/// Relation planner that intercepts `TABLESAMPLE` clauses in SQL and creates
/// [`TableSamplePlanNode`] logical nodes.
#[derive(Debug)]
struct TableSamplePlanner;

impl RelationPlanner for TableSamplePlanner {
    fn plan_relation(
        &self,
        relation: TableFactor,
        context: &mut dyn RelationPlannerContext,
    ) -> Result<RelationPlanning> {
        // Only handle Table relations with TABLESAMPLE clause
        let TableFactor::Table {
            sample: Some(sample),
            alias,
            name,
            args,
            with_hints,
            version,
            with_ordinality,
            partitions,
            json_path,
            index_hints,
        } = relation
        else {
            return Ok(RelationPlanning::Original(Box::new(relation)));
        };

        // Extract sample spec (handles both before/after alias positions)
        let sample = match sample {
            ast::TableSampleKind::BeforeTableAlias(s)
            | ast::TableSampleKind::AfterTableAlias(s) => s,
        };

        // Validate sampling method
        if let Some(method) = &sample.name
            && *method != TableSampleMethod::Bernoulli
            && *method != TableSampleMethod::Row
        {
            return not_impl_err!(
                "Sampling method {} is not supported (only BERNOULLI and ROW)",
                method
            );
        }

        // Offset sampling (ClickHouse-style) not supported
        if sample.offset.is_some() {
            return not_impl_err!(
                "TABLESAMPLE with OFFSET is not supported (requires total row count)"
            );
        }

        // Parse optional REPEATABLE seed
        let seed = sample
            .seed
            .map(|s| {
                s.value.to_string().parse::<u64>().map_err(|_| {
                    plan_datafusion_err!("REPEATABLE seed must be an integer")
                })
            })
            .transpose()?;

        // Plan the underlying table without the sample clause
        let base_relation = TableFactor::Table {
            sample: None,
            alias: alias.clone(),
            name,
            args,
            with_hints,
            version,
            with_ordinality,
            partitions,
            json_path,
            index_hints,
        };
        let input = context.plan(base_relation)?;

        // Handle bucket sampling (Hive-style: TABLESAMPLE(BUCKET x OUT OF y))
        if let Some(bucket) = sample.bucket {
            if bucket.on.is_some() {
                return not_impl_err!(
                    "TABLESAMPLE BUCKET with ON clause requires CLUSTERED BY table"
                );
            }
            let bucket_num: u64 =
                bucket.bucket.to_string().parse().map_err(|_| {
                    plan_datafusion_err!("bucket number must be an integer")
                })?;
            let total: u64 =
                bucket.total.to_string().parse().map_err(|_| {
                    plan_datafusion_err!("bucket total must be an integer")
                })?;

            let fraction = bucket_num as f64 / total as f64;
            let plan = TableSamplePlanNode::new(input, fraction, seed).into_plan();
            return Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
                plan, alias,
            ))));
        }

        // Handle quantity-based sampling
        let Some(quantity) = sample.quantity else {
            return plan_err!(
                "TABLESAMPLE requires a quantity (percentage, fraction, or row count)"
            );
        };
        let quantity_value_expr = context.sql_to_expr(quantity.value, input.schema())?;

        match quantity.unit {
            // TABLESAMPLE (N ROWS) - exact row limit
            Some(TableSampleUnit::Rows) => {
                let rows: i64 = parse_literal::<Int64Type>(&quantity_value_expr)?;
                if rows < 0 {
                    return plan_err!("row count must be non-negative, got {}", rows);
                }
                let plan = LogicalPlanBuilder::from(input)
                    .limit(0, Some(rows as usize))?
                    .build()?;
                Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
                    plan, alias,
                ))))
            }

            // TABLESAMPLE (N PERCENT) - percentage sampling
            Some(TableSampleUnit::Percent) => {
                let percent: f64 = parse_literal::<Float64Type>(&quantity_value_expr)?;
                let fraction = percent / 100.0;
                let plan = TableSamplePlanNode::new(input, fraction, seed).into_plan();
                Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
                    plan, alias,
                ))))
            }

            // TABLESAMPLE (N) - fraction if <1.0, row limit if >=1.0
            None => {
                let value = parse_literal::<Float64Type>(&quantity_value_expr)?;
                if value < 0.0 {
                    return plan_err!("sample value must be non-negative, got {}", value);
                }
                let plan = if value >= 1.0 {
                    // Interpret as row limit
                    LogicalPlanBuilder::from(input)
                        .limit(0, Some(value as usize))?
                        .build()?
                } else {
                    // Interpret as fraction
                    TableSamplePlanNode::new(input, value, seed).into_plan()
                };
                Ok(RelationPlanning::Planned(Box::new(PlannedRelation::new(
                    plan, alias,
                ))))
            }
        }
    }
}

/// Custom logical plan node representing a TABLESAMPLE operation.
///
/// Stores sampling parameters (bounds, seed) and wraps the input plan.
/// Gets converted to [`SampleExec`] during physical planning.
#[derive(Debug, Clone, Hash, Eq, PartialEq, PartialOrd)]
struct TableSamplePlanNode {
    input: LogicalPlan,
    lower_bound: HashableF64,
    upper_bound: HashableF64,
    seed: u64,
}

impl TableSamplePlanNode {
    /// Create a new sampling node with the given fraction (0.0 to 1.0).
    fn new(input: LogicalPlan, fraction: f64, seed: Option<u64>) -> Self {
        Self {
            input,
            lower_bound: HashableF64(0.0),
            upper_bound: HashableF64(fraction),
            seed: seed.unwrap_or_else(rand::random),
        }
    }

    /// Wrap this node in a LogicalPlan::Extension.
    fn into_plan(self) -> LogicalPlan {
        LogicalPlan::Extension(Extension {
            node: Arc::new(self),
        })
    }
}

impl UserDefinedLogicalNodeCore for TableSamplePlanNode {
    fn name(&self) -> &str {
        "TableSample"
    }

    fn inputs(&self) -> Vec<&LogicalPlan> {
        vec![&self.input]
    }

    fn schema(&self) -> &DFSchemaRef {
        self.input.schema()
    }

    fn expressions(&self) -> Vec<Expr> {
        vec![]
    }

    fn fmt_for_explain(&self, f: &mut Formatter) -> fmt::Result {
        write!(
            f,
            "Sample: bounds=[{}, {}], seed={}",
            self.lower_bound.0, self.upper_bound.0, self.seed
        )
    }

    fn with_exprs_and_inputs(
        &self,
        _exprs: Vec<Expr>,
        mut inputs: Vec<LogicalPlan>,
    ) -> Result<Self> {
        Ok(Self {
            input: inputs.swap_remove(0),
            lower_bound: self.lower_bound,
            upper_bound: self.upper_bound,
            seed: self.seed,
        })
    }
}

/// Wrapper for f64 that implements Hash and Eq (required for LogicalPlan).
#[derive(Debug, Clone, Copy, PartialOrd)]
struct HashableF64(f64);

impl PartialEq for HashableF64 {
    fn eq(&self, other: &Self) -> bool {
        self.0.to_bits() == other.0.to_bits()
    }
}

impl Eq for HashableF64 {}

impl Hash for HashableF64 {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.to_bits().hash(state);
    }
}

// ============================================================================
// Physical Planning: TableSampleQueryPlanner + TableSampleExtensionPlanner
// ============================================================================

/// Custom query planner that registers [`TableSampleExtensionPlanner`] to
/// convert [`TableSamplePlanNode`] into [`SampleExec`].
#[derive(Debug)]
struct TableSampleQueryPlanner;

#[async_trait]
impl QueryPlanner for TableSampleQueryPlanner {
    async fn create_physical_plan(
        &self,
        logical_plan: &LogicalPlan,
        session_state: &SessionState,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(
            TableSampleExtensionPlanner,
        )]);
        planner
            .create_physical_plan(logical_plan, session_state)
            .await
    }
}

/// Extension planner that converts [`TableSamplePlanNode`] to [`SampleExec`].
struct TableSampleExtensionPlanner;

#[async_trait]
impl ExtensionPlanner for TableSampleExtensionPlanner {
    async fn plan_extension(
        &self,
        _planner: &dyn PhysicalPlanner,
        node: &dyn UserDefinedLogicalNode,
        _logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
        _session_state: &SessionState,
    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
        let Some(sample_node) = node.as_any().downcast_ref::<TableSamplePlanNode>()
        else {
            return Ok(None);
        };

        let exec = SampleExec::try_new(
            Arc::clone(&physical_inputs[0]),
            sample_node.lower_bound.0,
            sample_node.upper_bound.0,
            sample_node.seed,
        )?;
        Ok(Some(Arc::new(exec)))
    }
}

// ============================================================================
// Physical Execution: SampleExec + BernoulliSampler
// ============================================================================

/// Physical execution plan that samples rows from its input using Bernoulli sampling.
///
/// Each row is independently selected with probability `(upper_bound - lower_bound)`
/// and appears at most once.
#[derive(Debug, Clone)]
pub struct SampleExec {
    input: Arc<dyn ExecutionPlan>,
    lower_bound: f64,
    upper_bound: f64,
    seed: u64,
    metrics: ExecutionPlanMetricsSet,
    cache: Arc<PlanProperties>,
}

impl SampleExec {
    /// Create a new SampleExec with Bernoulli sampling (without replacement).
    ///
    /// # Arguments
    /// * `input` - The input execution plan
    /// * `lower_bound` - Lower bound of sampling range (typically 0.0)
    /// * `upper_bound` - Upper bound of sampling range (0.0 to 1.0)
    /// * `seed` - Random seed for reproducible sampling
    pub fn try_new(
        input: Arc<dyn ExecutionPlan>,
        lower_bound: f64,
        upper_bound: f64,
        seed: u64,
    ) -> Result<Self> {
        if lower_bound < 0.0 || upper_bound > 1.0 || lower_bound > upper_bound {
            return internal_err!(
                "Sampling bounds must satisfy 0.0 <= lower <= upper <= 1.0, got [{}, {}]",
                lower_bound,
                upper_bound
            );
        }

        let cache = PlanProperties::new(
            EquivalenceProperties::new(input.schema()),
            input.properties().partitioning.clone(),
            input.properties().emission_type,
            input.properties().boundedness,
        );

        Ok(Self {
            input,
            lower_bound,
            upper_bound,
            seed,
            metrics: ExecutionPlanMetricsSet::new(),
            cache: Arc::new(cache),
        })
    }

    /// Create a sampler for the given partition.
    fn create_sampler(&self, partition: usize) -> BernoulliSampler {
        let seed = self.seed.wrapping_add(partition as u64);
        BernoulliSampler::new(self.lower_bound, self.upper_bound, seed)
    }
}

impl DisplayAs for SampleExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
        write!(
            f,
            "SampleExec: bounds=[{}, {}], seed={}",
            self.lower_bound, self.upper_bound, self.seed
        )
    }
}

impl ExecutionPlan for SampleExec {
    fn name(&self) -> &'static str {
        "SampleExec"
    }

    fn as_any(&self) -> &dyn Any {
        self
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        &self.cache
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        // Sampling preserves row order (rows are filtered, not reordered)
        vec![true]
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn with_new_children(
        self: Arc<Self>,
        mut children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        Ok(Arc::new(Self::try_new(
            children.swap_remove(0),
            self.lower_bound,
            self.upper_bound,
            self.seed,
        )?))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        Ok(Box::pin(SampleStream {
            input: self.input.execute(partition, context)?,
            sampler: self.create_sampler(partition),
            metrics: BaselineMetrics::new(&self.metrics, partition),
        }))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }

    fn partition_statistics(&self, partition: Option<usize>) -> Result<Statistics> {
        let mut stats = self.input.partition_statistics(partition)?;
        let ratio = self.upper_bound - self.lower_bound;

        // Scale statistics by sampling ratio (inexact due to randomness)
        stats.num_rows = stats
            .num_rows
            .map(|n| (n as f64 * ratio) as usize)
            .to_inexact();
        stats.total_byte_size = stats
            .total_byte_size
            .map(|n| (n as f64 * ratio) as usize)
            .to_inexact();

        Ok(stats)
    }
}

/// Bernoulli sampler: includes each row with probability `(upper - lower)`.
/// This is sampling **without replacement** - each row appears at most once.
struct BernoulliSampler {
    lower_bound: f64,
    upper_bound: f64,
    rng: StdRng,
}

impl BernoulliSampler {
    fn new(lower_bound: f64, upper_bound: f64, seed: u64) -> Self {
        Self {
            lower_bound,
            upper_bound,
            rng: StdRng::seed_from_u64(seed),
        }
    }

    fn sample(&mut self, batch: &RecordBatch) -> Result<RecordBatch> {
        let range = self.upper_bound - self.lower_bound;
        if range <= 0.0 {
            return Ok(RecordBatch::new_empty(batch.schema()));
        }

        // Select rows where random value falls in [lower, upper)
        let indices: Vec<u32> = (0..batch.num_rows())
            .filter(|_| {
                let r: f64 = self.rng.random();
                r >= self.lower_bound && r < self.upper_bound
            })
            .map(|i| i as u32)
            .collect();

        if indices.is_empty() {
            return Ok(RecordBatch::new_empty(batch.schema()));
        }

        compute::take_record_batch(batch, &UInt32Array::from(indices))
            .map_err(DataFusionError::from)
    }
}

/// Stream adapter that applies sampling to each batch.
struct SampleStream {
    input: SendableRecordBatchStream,
    sampler: BernoulliSampler,
    metrics: BaselineMetrics,
}

impl Stream for SampleStream {
    type Item = Result<RecordBatch>;

    fn poll_next(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        match ready!(self.input.poll_next_unpin(cx)) {
            Some(Ok(batch)) => {
                let elapsed = self.metrics.elapsed_compute().clone();
                let _timer = elapsed.timer();
                let result = self.sampler.sample(&batch);
                Poll::Ready(Some(result.record_output(&self.metrics)))
            }
            Some(Err(e)) => Poll::Ready(Some(Err(e))),
            None => Poll::Ready(None),
        }
    }
}

impl RecordBatchStream for SampleStream {
    fn schema(&self) -> SchemaRef {
        self.input.schema()
    }
}