indexlake-datafusion 0.6.0

use std::ops::Range;
use std::sync::Arc;

use arrow::array::{RecordBatch, RecordBatchOptions};
use arrow::datatypes::{Schema, SchemaRef};
use datafusion_common::stats::Precision;
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{DFSchema, DataFusionError, Statistics, project_schema};
use datafusion_execution::{SendableRecordBatchStream, TaskContext};
use datafusion_expr::Expr;
use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr};
use datafusion_physical_plan::display::ProjectSchemaDisplay;
use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use datafusion_physical_plan::{
    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties,
};
use futures::{StreamExt, TryStreamExt};
use indexlake::catalog::DataFileRecord;
use indexlake::table::{Table, TableScan, TableScanPartition};
use log::error;

use crate::{LazyTable, datafusion_expr_to_indexlake_expr, schema_projection_equals};

#[derive(Debug)]
pub struct IndexLakeScanExec {
    pub lazy_table: LazyTable,
    pub output_schema: SchemaRef,
    pub partition_count: usize,
    scan_partitions: Arc<Vec<TableScanPartition>>,
    partition_row_counts: Arc<Vec<usize>>,
    pub projection: Option<Vec<usize>>,
    pub filters: Vec<Expr>,
    pub batch_size: usize,
    pub limit: Option<usize>,
    properties: Arc<PlanProperties>,
}

impl IndexLakeScanExec {
    #[allow(clippy::too_many_arguments)]
    pub fn try_new(
        lazy_table: LazyTable,
        output_schema: SchemaRef,
        scan_partitions: Arc<Vec<TableScanPartition>>,
        partition_row_counts: Arc<Vec<usize>>,
        projection: Option<Vec<usize>>,
        filters: Vec<Expr>,
        batch_size: usize,
        limit: Option<usize>,
    ) -> Result<Self, DataFusionError> {
        let partition_count = scan_partitions.len();
        if partition_row_counts.len() != partition_count {
            return Err(DataFusionError::Plan(format!(
                "partition row count mismatch: {partition_count} partitions, {} row counts",
                partition_row_counts.len()
            )));
        }
        let projected_schema = project_schema(&output_schema, projection.as_ref())?;
        let properties = Arc::new(PlanProperties::new(
            EquivalenceProperties::new(projected_schema),
            Partitioning::UnknownPartitioning(partition_count),
            EmissionType::Incremental,
            Boundedness::Bounded,
        ));
        Ok(Self {
            lazy_table,
            output_schema,
            partition_count,
            scan_partitions,
            partition_row_counts,
            projection,
            filters,
            batch_size,
            limit,
            properties,
        })
    }

    pub fn get_scan_partition(&self, partition: Option<usize>) -> TableScanPartition {
        match partition {
            Some(partition) => self.scan_partitions[partition].clone(),
            None => TableScanPartition::single_partition(),
        }
    }

    pub(crate) fn scan_partitions(&self) -> &Arc<Vec<TableScanPartition>> {
        &self.scan_partitions
    }

    pub(crate) fn partition_row_counts(&self) -> &Arc<Vec<usize>> {
        &self.partition_row_counts
    }
}

impl ExecutionPlan for IndexLakeScanExec {
    fn name(&self) -> &str {
        "IndexLakeScanExec"
    }

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

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

    fn with_new_children(
        self: Arc<Self>,
        _children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        Ok(self)
    }

    fn apply_expressions(
        &self,
        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion, DataFusionError>,
    ) -> Result<TreeNodeRecursion, DataFusionError> {
        // Filters are logical `datafusion_expr::Expr`, not physical expressions.
        Ok(TreeNodeRecursion::Continue)
    }

    fn execute(
        &self,
        partition: usize,
        _context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream, DataFusionError> {
        if partition >= self.partition_count {
            return Err(DataFusionError::Execution(format!(
                "partition index out of range: {partition} >= {}",
                self.partition_count
            )));
        }

        let df_schema = DFSchema::try_from(self.output_schema.as_ref().clone())?;
        let il_filters = self
            .filters
            .iter()
            .map(|f| datafusion_expr_to_indexlake_expr(f, &df_schema))
            .collect::<Result<Vec<_>, _>>()?;

        let scan_partition = self.get_scan_partition(Some(partition));

        let scan = TableScan::default()
            .with_projection(self.projection.clone())
            .with_filters(il_filters)
            .with_batch_size(self.batch_size)
            .with_partition(scan_partition)
            .with_limit(self.limit);

        let projected_schema = self.schema();
        let lazy_table = self.lazy_table.clone();

        let fut = async move {
            let table = lazy_table.get_or_load().await?;
            get_batch_stream(table, projected_schema.clone(), scan).await
        };
        let stream = futures::stream::once(fut).try_flatten();
        Ok(Box::pin(RecordBatchStreamAdapter::new(
            self.schema(),
            stream,
        )))
    }

    fn partition_statistics(
        &self,
        partition: Option<usize>,
    ) -> Result<Arc<Statistics>, DataFusionError> {
        let row_count = if let Some(partition) = partition {
            *self.partition_row_counts.get(partition).ok_or_else(|| {
                DataFusionError::Plan(format!(
                    "partition index out of range: {partition} >= {}",
                    self.partition_row_counts.len()
                ))
            })?
        } else {
            self.partition_row_counts.iter().copied().sum()
        };

        if self.filters.is_empty() {
            if let Some(limit) = self.limit {
                Ok(Arc::new(Statistics {
                    num_rows: Precision::Exact(std::cmp::min(row_count, limit)),
                    total_byte_size: Precision::Absent,
                    column_statistics: Statistics::unknown_column(&self.schema()),
                }))
            } else {
                Ok(Arc::new(Statistics {
                    num_rows: Precision::Exact(row_count),
                    total_byte_size: Precision::Absent,
                    column_statistics: Statistics::unknown_column(&self.schema()),
                }))
            }
        } else {
            Ok(Arc::new(Statistics {
                num_rows: Precision::Inexact(row_count),
                total_byte_size: Precision::Absent,
                column_statistics: Statistics::unknown_column(&self.schema()),
            }))
        }
    }

    fn with_fetch(&self, limit: Option<usize>) -> Option<Arc<dyn ExecutionPlan>> {
        match IndexLakeScanExec::try_new(
            self.lazy_table.clone(),
            self.output_schema.clone(),
            self.scan_partitions.clone(),
            self.partition_row_counts.clone(),
            self.projection.clone(),
            self.filters.clone(),
            self.batch_size,
            limit,
        ) {
            Ok(exec) => Some(Arc::new(exec)),
            Err(e) => {
                error!("[indexlake] Failed to create IndexLakeScanExec with fetch: {e}");
                None
            }
        }
    }

    fn fetch(&self) -> Option<usize> {
        self.limit
    }
}

impl DisplayAs for IndexLakeScanExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "IndexLakeScanExec: table={}.{}, partitions={}",
            self.lazy_table.namespace_name, self.lazy_table.table_name, self.partition_count
        )?;
        let projected_schema = self.schema();
        if !schema_projection_equals(&projected_schema, &self.output_schema) {
            write!(
                f,
                ", projection={}",
                ProjectSchemaDisplay(&projected_schema)
            )?;
        }
        if !self.filters.is_empty() {
            write!(
                f,
                ", filters=[{}]",
                self.filters
                    .iter()
                    .map(|f| f.to_string())
                    .collect::<Vec<_>>()
                    .join(", ")
            )?;
        }
        if let Some(limit) = self.limit {
            write!(f, ", limit={limit}")?;
        }
        Ok(())
    }
}

pub(crate) fn build_scan_partitions(
    partition_count: usize,
    data_files: Option<Arc<Vec<DataFileRecord>>>,
) -> Vec<TableScanPartition> {
    if partition_count == 0 {
        return Vec::new();
    }

    match data_files {
        Some(data_files) => {
            let ranges = calc_data_file_partition_ranges(partition_count, data_files.len());
            (0..partition_count)
                .map(|partition| {
                    let range = ranges[partition].clone();
                    TableScanPartition::Provided {
                        contains_inline_rows: partition == 0,
                        data_file_records: if let Some(range) = range {
                            data_files[range].to_vec()
                        } else {
                            vec![]
                        },
                    }
                })
                .collect()
        }
        None => (0..partition_count)
            .map(|partition| TableScanPartition::Auto {
                partition_idx: partition,
                partition_count,
            })
            .collect(),
    }
}

async fn get_batch_stream(
    table: Arc<Table>,
    projected_schema: SchemaRef,
    mut scan: TableScan,
) -> Result<SendableRecordBatchStream, DataFusionError> {
    let stream = if scan.projection == Some(Vec::new()) {
        scan.projection = Some(vec![0]);
        let stream = table.scan(scan).await?;
        stream
            .map(|batch| {
                let batch = batch?;
                let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows()));
                let new_batch =
                    RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options)?;
                Ok(new_batch)
            })
            .boxed()
    } else {
        table.scan(scan).await?
    };
    let stream = stream.map_err(DataFusionError::from);
    Ok(Box::pin(RecordBatchStreamAdapter::new(
        projected_schema,
        stream,
    )))
}

fn calc_data_file_partition_ranges(
    partition_count: usize,
    data_file_count: usize,
) -> Vec<Option<Range<usize>>> {
    if partition_count == 0 {
        return vec![None];
    }

    // Each partition owns `base` files, and the first `extras` partitions each
    // take one extra. Partition i's range starts after i*base evenly-divided
    // files plus min(i, extras) extras already handed to earlier partitions,
    // so partition sizes differ by at most 1.
    let base = data_file_count / partition_count;
    let extras = data_file_count % partition_count;

    (0..partition_count)
        .map(|i| {
            let len = base + usize::from(i < extras);
            let start = i * base + i.min(extras);
            (len != 0).then(|| start..start + len)
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_partition_data_file_range() {
        let ranges = calc_data_file_partition_ranges(2, 0);
        assert_eq!(ranges, vec![None, None]);

        let ranges = calc_data_file_partition_ranges(2, 1);
        assert_eq!(ranges, vec![Some(0..1), None]);

        let ranges = calc_data_file_partition_ranges(2, 2);
        assert_eq!(ranges, vec![Some(0..1), Some(1..2)]);

        let ranges = calc_data_file_partition_ranges(2, 3);
        assert_eq!(ranges, vec![Some(0..2), Some(2..3)]);

        let ranges = calc_data_file_partition_ranges(2, 4);
        assert_eq!(ranges, vec![Some(0..2), Some(2..4)]);
    }
}