indexlake-datafusion 0.6.0

use std::collections::HashMap;
use std::sync::Arc;

use arrow::datatypes::SchemaRef;
use datafusion_catalog::{Session, TableProvider};
use datafusion_common::DataFusionError;
use datafusion_common::stats::Precision;
use datafusion_common::{DFSchema, Statistics};
use datafusion_expr::Expr;
use datafusion_expr::TableProviderFilterPushDown;
use datafusion_expr::TableType;
use datafusion_expr::dml::InsertOp;
use datafusion_physical_plan::ExecutionPlan;
use indexlake::Client;
use indexlake::index::FilterSupport;
use indexlake::table::{Table, TableScanPartition, TableUpdate};
use indexlake::utils::schema_without_row_id;
use tokio::sync::Mutex;

use crate::scan::build_scan_partitions;
use crate::{
    IndexLakeDeleteExec, IndexLakeInsertExec, IndexLakeScanExec, IndexLakeUpdateExec,
    datafusion_expr_to_indexlake_expr, indexlake_expr_to_datafusion_expr,
};

#[derive(Debug)]
pub struct IndexLakeTable {
    client: Arc<Client>,
    table: Arc<Table>,
    batch_size: usize,
    num_scan_partitions: usize,
    column_defaults: HashMap<String, Expr>,
    hide_row_id: bool,
    bypass_insert_threshold: usize,
    table_row_count: Option<usize>,
}

impl IndexLakeTable {
    pub fn try_new(client: Arc<Client>, table: Arc<Table>) -> Result<Self, DataFusionError> {
        let mut column_defaults = HashMap::new();
        for field_record in table.field_records.iter() {
            if let Some(default_value) = &field_record.default_value {
                let expr = indexlake_expr_to_datafusion_expr(default_value)?;
                column_defaults.insert(field_record.field_name.clone(), expr);
            }
        }
        Ok(Self {
            client,
            table,
            batch_size: 2048,
            num_scan_partitions: 16,
            column_defaults,
            hide_row_id: false,
            bypass_insert_threshold: 1000,
            table_row_count: None,
        })
    }

    pub async fn with_table_statistics(mut self) -> Result<Self, DataFusionError> {
        let counts = self
            .table
            .count(&[TableScanPartition::single_partition()])
            .await?;
        self.table_row_count = Some(counts[0]);
        Ok(self)
    }

    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    pub fn with_num_scan_partitions(mut self, num_scan_partitions: usize) -> Self {
        self.num_scan_partitions = num_scan_partitions;
        self
    }

    pub fn with_hide_row_id(mut self, hide_row_id: bool) -> Self {
        self.hide_row_id = hide_row_id;
        self
    }

    pub fn with_bypass_insert_threshold(mut self, bypass_insert_threshold: usize) -> Self {
        self.bypass_insert_threshold = bypass_insert_threshold;
        self
    }
}

#[async_trait::async_trait]
impl TableProvider for IndexLakeTable {
    fn schema(&self) -> SchemaRef {
        if self.hide_row_id {
            Arc::new(schema_without_row_id(&self.table.output_schema))
        } else {
            self.table.output_schema.clone()
        }
    }

    fn table_type(&self) -> TableType {
        TableType::Base
    }

    fn get_column_default(&self, column: &str) -> Option<&Expr> {
        self.column_defaults.get(column)
    }

    async fn scan(
        &self,
        _state: &dyn Session,
        projection: Option<&Vec<usize>>,
        filters: &[Expr],
        limit: Option<usize>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let data_file_count = self.table.data_file_count().await?;
        let data_files = if data_file_count > 1000 {
            None
        } else {
            let records = self.table.data_file_records().await?;
            Some(Arc::new(records))
        };

        let il_projection = if let Some(df_projection) = projection
            && self.hide_row_id
        {
            Some(df_projection.iter().map(|i| i + 1).collect::<Vec<_>>())
        } else {
            projection.cloned()
        };

        let lazy_table = LazyTable::new(
            self.client.clone(),
            self.table.namespace_name.clone(),
            self.table.table_name.clone(),
        )
        .with_table(self.table.clone());

        let scan_partitions = Arc::new(build_scan_partitions(self.num_scan_partitions, data_files));
        let partition_row_counts = self.table.count(scan_partitions.as_ref()).await?;
        let partition_row_counts = Arc::new(partition_row_counts);
        let exec = IndexLakeScanExec::try_new(
            lazy_table,
            self.table.output_schema.clone(),
            scan_partitions,
            partition_row_counts,
            il_projection,
            filters.to_vec(),
            self.batch_size,
            limit,
        )?;
        Ok(Arc::new(exec))
    }

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> Result<Vec<TableProviderFilterPushDown>, DataFusionError> {
        let df_schema = DFSchema::try_from(self.table.output_schema.clone())?;
        let mut supports = Vec::with_capacity(filters.len());
        for filter in filters {
            let Ok(il_expr) = datafusion_expr_to_indexlake_expr(filter, &df_schema) else {
                supports.push(TableProviderFilterPushDown::Unsupported);
                continue;
            };
            let support = self.table.supports_filter(il_expr.clone())?;
            match support {
                FilterSupport::Exact => supports.push(TableProviderFilterPushDown::Exact),
                FilterSupport::Inexact => supports.push(TableProviderFilterPushDown::Inexact),
                FilterSupport::Unsupported => {
                    supports.push(TableProviderFilterPushDown::Unsupported)
                }
            }
        }
        Ok(supports)
    }

    fn statistics(&self) -> Option<Statistics> {
        let row_count = self.table_row_count?;
        Some(Statistics {
            num_rows: Precision::Exact(row_count),
            total_byte_size: Precision::Absent,
            column_statistics: Statistics::unknown_column(&self.table.output_schema),
        })
    }

    async fn insert_into(
        &self,
        _state: &dyn Session,
        input: Arc<dyn ExecutionPlan>,
        insert_op: InsertOp,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let lazy_table = LazyTable::new(
            self.client.clone(),
            self.table.namespace_name.clone(),
            self.table.table_name.clone(),
        )
        .with_table(self.table.clone());

        let insert_exec = IndexLakeInsertExec::try_new(
            lazy_table,
            input,
            insert_op,
            self.bypass_insert_threshold,
        )?;

        Ok(Arc::new(insert_exec))
    }

    async fn delete_from(
        &self,
        _state: &dyn Session,
        filters: Vec<Expr>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let df_schema = DFSchema::try_from(self.table.output_schema.clone())?;

        let mut condition = indexlake::expr::lit(true);
        for filter in filters {
            let il_filter = datafusion_expr_to_indexlake_expr(&filter, &df_schema)?;
            condition = indexlake::expr::Expr::BinaryExpr(indexlake::expr::BinaryExpr {
                left: Box::new(condition),
                op: indexlake::expr::BinaryOp::And,
                right: Box::new(il_filter),
            });
        }

        let lazy_table = LazyTable::new(
            self.client.clone(),
            self.table.namespace_name.clone(),
            self.table.table_name.clone(),
        )
        .with_table(self.table.clone());

        let exec = IndexLakeDeleteExec::try_new(lazy_table, condition)?;
        Ok(Arc::new(exec))
    }

    async fn update(
        &self,
        _state: &dyn Session,
        assignments: Vec<(String, Expr)>,
        filters: Vec<Expr>,
    ) -> Result<Arc<dyn ExecutionPlan>, DataFusionError> {
        let df_schema = DFSchema::try_from(self.table.output_schema.clone())?;

        let mut condition = indexlake::expr::lit(true);
        for filter in filters {
            let il_filter = datafusion_expr_to_indexlake_expr(&filter, &df_schema)?;
            condition = indexlake::expr::Expr::BinaryExpr(indexlake::expr::BinaryExpr {
                left: Box::new(condition),
                op: indexlake::expr::BinaryOp::And,
                right: Box::new(il_filter),
            });
        }

        let mut set_map = HashMap::with_capacity(assignments.len());
        for (col, expr) in assignments {
            let il_expr = datafusion_expr_to_indexlake_expr(&expr, &df_schema)?;
            set_map.insert(col, il_expr);
        }

        let update = TableUpdate { set_map, condition };

        let lazy_table = LazyTable::new(
            self.client.clone(),
            self.table.namespace_name.clone(),
            self.table.table_name.clone(),
        )
        .with_table(self.table.clone());

        let exec = IndexLakeUpdateExec::try_new(lazy_table, update)?;
        Ok(Arc::new(exec))
    }
}

/// A lazy-loaded table holder containing client and table metadata.
/// The actual Table is loaded on first access if not provided upfront.
#[derive(Debug, Clone)]
pub struct LazyTable {
    pub client: Arc<Client>,
    pub namespace_name: String,
    pub table_name: String,
    table: Arc<Mutex<Option<Arc<Table>>>>,
}

impl LazyTable {
    /// Create a new LazyTable without a pre-loaded table.
    pub fn new(client: Arc<Client>, namespace_name: String, table_name: String) -> Self {
        Self {
            client,
            namespace_name,
            table_name,
            table: Arc::new(Mutex::new(None)),
        }
    }

    /// Create a new LazyTable with a pre-loaded table.
    pub fn with_table(mut self, table: Arc<Table>) -> Self {
        self.table = Arc::new(Mutex::new(Some(table)));
        self
    }

    /// Get the table, loading it lazily if not already loaded.
    pub async fn get_or_load(&self) -> Result<Arc<Table>, indexlake::ILError> {
        let mut guard = self.table.lock().await;
        if let Some(table) = guard.as_ref() {
            return Ok(table.clone());
        }

        let table = self
            .client
            .load_table(&self.namespace_name, &self.table_name)
            .await?;
        let table = Arc::new(table);
        *guard = Some(table.clone());
        Ok(table)
    }
}