timeseries_table_format/datafusion/

ts_table_provider.rs

mod segment_pruning;
#[cfg(test)]
mod tests;
mod timestamp_pruning;

use crate::storage::file_size;

use std::collections::HashSet;
use std::path::Path;
use std::sync::Arc;

use arrow::datatypes::{DataType, SchemaRef};
use async_trait::async_trait;

use datafusion::catalog::Session;
use datafusion::catalog::TableProvider;
use datafusion::common::DFSchema;

use datafusion::datasource::listing::PartitionedFile;
use datafusion::datasource::physical_plan::FileScanConfigBuilder;
use datafusion::datasource::physical_plan::ParquetSource;
use datafusion::datasource::source::DataSourceExec;
use datafusion::error::{DataFusionError, Result as DFResult};
use datafusion::execution::object_store::ObjectStoreUrl;

use datafusion::logical_expr::{Expr, Operator};

use datafusion::logical_expr::TableProviderFilterPushDown;

use crate::metadata::index::{IndexKind, IndexSpec};
use crate::table::TimeSeriesTable;
use crate::transaction_log::SegmentMeta;
use crate::transaction_log::TableState;
use datafusion::logical_expr::utils::{conjunction, expr_to_columns};
use datafusion::physical_expr::PhysicalExpr;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::expressions::lit;
use datafusion::scalar::ScalarValue;
use tokio::sync::RwLock;

/// DataFusion table provider for a timeseries table schema.
///
/// The schema is captured when the provider is constructed. If the table schema
/// evolves, re-register a new provider to pick up the updated schema.
#[derive(Debug)]
pub struct TsTableProvider {
    table: Arc<TimeSeriesTable>,
    schema: SchemaRef,
    cache: RwLock<Cache>,

    object_store_url: ObjectStoreUrl,
}

#[derive(Debug)]
struct Cache {
    version: Option<u64>,
    state: Option<TableState>,
}

/// Wrap a generic error for DataFusion APIs.
fn df_external<E>(e: E) -> DataFusionError
where
    E: std::error::Error + Send + Sync + 'static,
{
    DataFusionError::External(Box::new(e))
}

fn metadata_pruning_expr(
    expr: &Expr,
    index: &IndexSpec,
    schema: &SchemaRef,
) -> DFResult<Option<Expr>> {
    if is_exact_entity_equality(expr, index, schema) {
        return Ok(Some(expr.clone()));
    }

    if let Expr::BinaryExpr(binary) = expr
        && matches!(binary.op, Operator::And | Operator::Or)
    {
        let left = metadata_pruning_expr(&binary.left, index, schema)?;
        let right = metadata_pruning_expr(&binary.right, index, schema)?;
        return Ok(if binary.op == Operator::And {
            match (left, right) {
                (Some(left), Some(right)) => Some(left.and(right)),
                (left @ Some(_), None) | (None, left @ Some(_)) => left,
                (None, None) => None,
            }
        } else {
            left.zip(right).map(|(left, right)| left.or(right))
        });
    }

    let mut columns = HashSet::new();
    expr_to_columns(expr, &mut columns)?;
    Ok(
        (!columns.is_empty() && columns.iter().all(|column| column.name == index.column))
            .then(|| expr.clone()),
    )
}

fn is_exact_entity_equality(expr: &Expr, index: &IndexSpec, schema: &SchemaRef) -> bool {
    let Expr::BinaryExpr(binary) = expr else {
        return false;
    };
    if binary.op != Operator::Eq {
        return false;
    }
    let (column, literal) = match (&*binary.left, &*binary.right) {
        (Expr::Column(column), Expr::Literal(literal, _))
        | (Expr::Literal(literal, _), Expr::Column(column)) => (column, literal),
        _ => return false,
    };
    entity_literal_matches(column.name.as_str(), literal, index, schema)
}

fn entity_literal_matches(
    column: &str,
    literal: &ScalarValue,
    index: &IndexSpec,
    schema: &SchemaRef,
) -> bool {
    if !index.entity_columns.iter().any(|entity| entity == column) {
        return false;
    }
    let Ok(field) = schema.field_with_name(column) else {
        return false;
    };
    matches!(
        (field.data_type(), literal),
        (DataType::Utf8, ScalarValue::Utf8(Some(_)))
            | (DataType::LargeUtf8, ScalarValue::LargeUtf8(Some(_)))
            | (DataType::Int32, ScalarValue::Int32(Some(_)))
            | (DataType::Int64, ScalarValue::Int64(Some(_)))
            | (DataType::UInt64, ScalarValue::UInt64(Some(_)))
    )
}

impl TsTableProvider {
    /// Creates a new provider backed by the given `TimeSeriesTable`.
    pub fn try_new(table: Arc<TimeSeriesTable>) -> DFResult<Self> {
        // Use the table's current in-memory snapshot to get schema.
        // (No schema evolution in v0.1, so this is stable.)
        let schema = table
            .state()
            .table_meta
            .arrow_schema_ref()
            .map_err(df_external)?;

        let object_store_url =
            ObjectStoreUrl::parse(table.location().object_store_url()).map_err(df_external)?;
        let state = table.state().clone();

        Ok(Self {
            table,
            schema,
            cache: RwLock::new(Cache {
                version: Some(state.version),
                state: Some(state),
            }),
            object_store_url,
        })
    }

    async fn latest_state(&self) -> DFResult<TableState> {
        let current_version = self.table.current_version().await.map_err(df_external)?;

        // Fast path: cache hit
        {
            let cache = self.cache.read().await;
            if cache.version == Some(current_version)
                && let Some(st) = cache.state.clone()
            {
                return Ok(st);
            }
        }

        // Refresh from log
        let state = self.table.load_latest_state().await.map_err(df_external)?;
        let mut cache = self.cache.write().await;
        cache.version = Some(state.version);
        cache.state = Some(state.clone());
        Ok(state)
    }

    async fn segment_file_size(&self, seg: &SegmentMeta) -> datafusion::error::Result<u64> {
        if let Some(sz) = seg.file_size {
            return Ok(sz);
        }

        let sz = file_size(self.table.location().storage(), Path::new(&seg.path))
            .await
            .map_err(|source| {
                df_external(source).context(format!(
                    "missing segment file size; failed to inspect {}",
                    seg.path
                ))
            })?;
        Ok(sz)
    }

    /// Return the ordered-index column name from the table's index spec.
    fn index_column_name(&self) -> &str {
        self.table.index_spec().column.as_str()
    }

    fn prune_segments_by_metadata<'a>(
        &self,
        segments: Vec<&'a SegmentMeta>,
        metadata_filters: &[Expr],
        pruning_predicate: &Arc<dyn PhysicalExpr>,
    ) -> DFResult<Vec<&'a SegmentMeta>> {
        if metadata_filters.is_empty() {
            return Ok(segments);
        }

        segment_pruning::prune_segments(
            &self.schema,
            self.table.index_spec(),
            segments,
            pruning_predicate,
        )
    }
}

#[async_trait]
impl TableProvider for TsTableProvider {
    fn schema(&self) -> arrow::datatypes::SchemaRef {
        Arc::clone(&self.schema)
    }

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

    fn supports_filters_pushdown(
        &self,
        filters: &[&Expr],
    ) -> DFResult<Vec<TableProviderFilterPushDown>> {
        // Inexact: we may prune files, and Parquet may prune row groups/pages,
        // but DataFusion will still apply the filter for correctness.
        Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()])
    }

    #[tracing::instrument(
        name = "table.scan.plan",
        level = "debug",
        skip_all,
        fields(
            snapshot_version = tracing::field::Empty,
            total_candidate_segments = tracing::field::Empty,
            selected_segments = tracing::field::Empty,
            pruned_segments = tracing::field::Empty,
            filter_count = filters.len(),
            projection_column_count = tracing::field::Empty,
            limit = tracing::field::Empty
        )
    )]
    async fn scan(
        &self,
        state: &dyn Session,
        projection: Option<&Vec<usize>>,
        filters: &[Expr], // may include all WHERE predicates
        limit: Option<usize>,
    ) -> DFResult<Arc<dyn ExecutionPlan>> {
        let span = tracing::Span::current();
        if let Some(projection) = projection {
            span.record("projection_column_count", projection.len());
        }
        if let Some(limit) = limit {
            span.record("limit", limit);
        }

        // 1) Get a snapshot (TableState) from core table
        let snapshot = self.latest_state().await?;
        span.record("snapshot_version", snapshot.version);

        for segment in snapshot.segments.values() {
            segment
                .validate_bounds(&self.table.index_spec().kind)
                .map_err(df_external)?;
        }

        let segments = snapshot.segments_sorted_by_index().map_err(df_external)?;
        let total_candidate_segments = segments.len();
        span.record("total_candidate_segments", total_candidate_segments);

        let df_schema = DFSchema::try_from(self.schema().as_ref().clone())?;
        let exact_predicate = conjunction(filters.to_vec())
            .map(|p| state.create_physical_expr(p, &df_schema))
            .transpose()?
            .unwrap_or_else(|| lit(true));

        let metadata_filters = filters
            .iter()
            .map(|filter| metadata_pruning_expr(filter, self.table.index_spec(), &self.schema))
            .collect::<DFResult<Vec<_>>>()?
            .into_iter()
            .flatten()
            .collect::<Vec<_>>();

        let pruning_filters = if matches!(self.table.index_spec().kind, IndexKind::Timestamp { .. })
        {
            let index_column = self.index_column_name();
            let index_type = self.schema.field_with_name(index_column)?.data_type();
            metadata_filters
                .iter()
                .cloned()
                .map(|filter| {
                    timestamp_pruning::normalize_timestamp_predicate(
                        filter,
                        index_column,
                        index_type,
                    )
                })
                .collect::<DFResult<Vec<_>>>()?
        } else {
            metadata_filters.clone()
        };
        let pruning_predicate = if pruning_filters.as_slice() == filters {
            Arc::clone(&exact_predicate)
        } else {
            conjunction(pruning_filters)
                .map(|p| state.create_physical_expr(p, &df_schema))
                .transpose()?
                .unwrap_or_else(|| lit(true))
        };

        // Build Parquet scan plan (DataSourceExec + ParquetSource)
        let parquet_source = Arc::new(
            ParquetSource::new(Arc::clone(&self.schema))
                .with_predicate(Arc::clone(&exact_predicate)),
        );

        let mut builder = FileScanConfigBuilder::new(self.object_store_url.clone(), parquet_source)
            .with_projection_indices(projection.cloned())?
            .with_limit(limit);

        let selected =
            self.prune_segments_by_metadata(segments, &metadata_filters, &pruning_predicate)?;
        span.record("selected_segments", selected.len());
        span.record(
            "pruned_segments",
            total_candidate_segments.saturating_sub(selected.len()),
        );
        for seg in selected {
            let file_size = self.segment_file_size(seg).await?;
            let location = self
                .table
                .location()
                .object_store_path(Path::new(&seg.path))
                .map_err(df_external)?;
            let pf = PartitionedFile::new(location.as_ref(), file_size);

            builder = builder.with_file(pf);
        }

        // Produce the execution plan
        let plan = DataSourceExec::from_data_source(builder.build());
        Ok(plan)
    }
}