datafusion-ducklake 0.7.0

//! User-Defined Table Functions (UDTFs) for DuckLake catalog metadata

use datafusion::catalog::TableFunctionImpl;
use datafusion::common::{Result as DataFusionResult, ScalarValue, plan_err};
use datafusion::datasource::TableProvider;
use datafusion::logical_expr::Expr;
use std::sync::Arc;

use crate::information_schema::{FilesTable, SnapshotsTable, TableInfoTable};
use crate::metadata_provider::{
    MetadataProvider, parse_snapshot_timestamp, require_snapshot, resolve_snapshot_at_or_after,
    resolve_snapshot_at_or_before,
};
use crate::path_resolver::{parse_object_store_url, resolve_path};
use crate::table::DuckLakeTable;
use crate::table_changes::{TableChangesTable, TableInsertionsTable};
use crate::table_deletions::TableDeletionsTable;
use crate::types::build_arrow_schema;

#[derive(Debug)]
pub struct DucklakeSnapshotsFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeSnapshotsFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeSnapshotsFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        if !exprs.is_empty() {
            return plan_err!("ducklake_snapshots() takes no arguments");
        }

        Ok(Arc::new(SnapshotsTable::new(self.provider.clone())))
    }
}

#[derive(Debug)]
pub struct DucklakeTableInfoFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeTableInfoFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeTableInfoFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        if !exprs.is_empty() {
            return plan_err!("ducklake_table_info() takes no arguments");
        }

        Ok(Arc::new(TableInfoTable::new(self.provider.clone())))
    }
}

#[derive(Debug)]
pub struct DucklakeListFilesFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeListFilesFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeListFilesFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        if !exprs.is_empty() {
            return plan_err!("ducklake_list_files() takes no arguments");
        }

        Ok(Arc::new(FilesTable::new(self.provider.clone())))
    }
}

#[derive(Debug)]
pub struct DucklakeTableAtFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeTableAtFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeTableAtFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        const FN_NAME: &str = "ducklake_table_at";
        if exprs.len() != 3 {
            return plan_err!(
                "{FN_NAME}() requires 3 arguments: \
                 {FN_NAME}('schema', 'table', version_or_timestamp)"
            );
        }
        let string_arg = |expr: &Expr, name: &str| match expr {
            Expr::Literal(ScalarValue::Utf8(Some(value)), _) => Ok(value.clone()),
            _ => plan_err!("{name} of {FN_NAME}() must be a string literal"),
        };
        let schema_name = string_arg(&exprs[0], "schema")?;
        let table_name = string_arg(&exprs[1], "table")?;
        let snapshot_id = resolve_snapshot_bound(
            &self.provider,
            &exprs[2],
            FN_NAME,
            "version_or_timestamp",
            false,
        )?;
        require_snapshot(self.provider.as_ref(), snapshot_id)
            .map_err(|error| datafusion::error::DataFusionError::Plan(error.to_string()))?;

        let schema = self
            .provider
            .get_schema_by_name(&schema_name, snapshot_id)
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Plan(format!(
                    "Schema '{schema_name}' does not exist at snapshot {snapshot_id}"
                ))
            })?;
        let table = self
            .provider
            .get_table_by_name(schema.schema_id, &table_name, snapshot_id)
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?
            .ok_or_else(|| {
                datafusion::error::DataFusionError::Plan(format!(
                    "Table '{schema_name}.{table_name}' does not exist at snapshot {snapshot_id}"
                ))
            })?;
        let data_path = self
            .provider
            .get_data_path()
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?;
        let (object_store_url, catalog_path) = parse_object_store_url(&data_path)
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?;
        let schema_path = resolve_path(&catalog_path, &schema.path, schema.path_is_relative)
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?;
        let table_path = resolve_path(&schema_path, &table.path, table.path_is_relative)
            .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?;
        let provider = DuckLakeTable::new(
            table.table_id,
            table.table_name,
            Arc::clone(&self.provider),
            snapshot_id,
            Arc::new(object_store_url),
            table_path,
        )
        .map_err(|error| datafusion::error::DataFusionError::External(Box::new(error)))?;

        Ok(Arc::new(provider))
    }
}

#[derive(Debug)]
pub struct DucklakeTableChangesFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeTableChangesFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

/// Everything a CDC table function needs about its target table.
struct CdcTableContext {
    table_id: i64,
    object_store_url: datafusion::execution::object_store::ObjectStoreUrl,
    table_path: String,
    table_schema: Arc<arrow::datatypes::Schema>,
    /// The columns the `table_schema` was built from, carrying the field ids each
    /// data file's columns are resolved by.
    columns: Vec<crate::metadata_provider::DuckLakeTableColumn>,
}

/// Split `'schema.table'` (defaulting to schema `main`).
fn parse_table_name(table_name: &str) -> (&str, &str) {
    if let Some(dot_pos) = table_name.find('.') {
        (&table_name[..dot_pos], &table_name[dot_pos + 1..])
    } else {
        ("main", table_name)
    }
}

/// Resolve a snapshot bound argument: an integer snapshot id, or a timestamp
/// (a string literal or a `TIMESTAMP '...'` cast) resolved against the
/// catalog's snapshot times: a START bound resolves to the FIRST snapshot
/// at-or-after the timestamp, an END bound to the LAST snapshot at-or-before
/// it.
fn resolve_snapshot_bound(
    provider: &Arc<dyn MetadataProvider>,
    expr: &Expr,
    fn_name: &str,
    arg_name: &str,
    is_start: bool,
) -> DataFusionResult<i64> {
    // `TIMESTAMP '...'` parses as a cast around a literal; unwrap it.
    let mut expr = expr;
    while let Expr::Cast(cast) = expr {
        expr = cast.expr.as_ref();
    }
    let ts: chrono::NaiveDateTime = match expr {
        Expr::Literal(ScalarValue::Int64(Some(v)), _) => return Ok(*v),
        Expr::Literal(ScalarValue::Int32(Some(v)), _) => return Ok(*v as i64),
        Expr::Literal(ScalarValue::Utf8(Some(s)), _) => match parse_snapshot_timestamp(s) {
            Some(ts) => ts,
            None => {
                return plan_err!(
                    "{arg_name} of {fn_name}() is not a valid timestamp: '{s}' \
                     (expected 'YYYY-MM-DD[ HH:MM:SS[.ffffff]]', UTC)"
                );
            },
        },
        Expr::Literal(other, _) => {
            let unit_and_value = match other {
                ScalarValue::TimestampSecond(Some(v), _) => Some((*v, 1_000_000_000i64)),
                ScalarValue::TimestampMillisecond(Some(v), _) => Some((*v, 1_000_000)),
                ScalarValue::TimestampMicrosecond(Some(v), _) => Some((*v, 1_000)),
                ScalarValue::TimestampNanosecond(Some(v), _) => Some((*v, 1)),
                _ => None,
            };
            match unit_and_value {
                Some((v, factor)) => {
                    let nanos = v.saturating_mul(factor);
                    match chrono::DateTime::from_timestamp(
                        nanos.div_euclid(1_000_000_000),
                        (nanos.rem_euclid(1_000_000_000)) as u32,
                    ) {
                        Some(dt) => dt.naive_utc(),
                        None => {
                            return plan_err!(
                                "{arg_name} of {fn_name}() is out of timestamp range"
                            );
                        },
                    }
                },
                None => {
                    return plan_err!(
                        "{arg_name} of {fn_name}() must be an integer snapshot id or a \
                         timestamp literal"
                    );
                },
            }
        },
        _ => {
            return plan_err!(
                "{arg_name} of {fn_name}() must be an integer snapshot id or a timestamp \
                 literal"
            );
        },
    };

    let resolved = if is_start {
        resolve_snapshot_at_or_after(provider.as_ref(), ts)
    } else {
        resolve_snapshot_at_or_before(provider.as_ref(), ts)
    };
    resolved
        .map_err(|error| datafusion::error::DataFusionError::Plan(format!("{fn_name}(): {error}")))
}

/// Parse the common `('schema.table', start, end)` argument list of the CDC
/// table functions and resolve the target table.
fn parse_cdc_args(
    provider: &Arc<dyn MetadataProvider>,
    exprs: &[Expr],
    fn_name: &str,
) -> DataFusionResult<(i64, i64, CdcTableContext)> {
    if exprs.len() != 3 {
        return plan_err!(
            "{fn_name}() requires 3 arguments: \
             {fn_name}('schema.table', start_snapshot, end_snapshot)"
        );
    }
    let table_name = match &exprs[0] {
        Expr::Literal(ScalarValue::Utf8(Some(name)), _) => name.clone(),
        _ => {
            return plan_err!(
                "First argument to {fn_name}() must be a string literal \
                 (e.g., 'main.users' or 'users')"
            );
        },
    };
    let start_snapshot =
        resolve_snapshot_bound(provider, &exprs[1], fn_name, "start_snapshot", true)?;
    let end_snapshot = resolve_snapshot_bound(provider, &exprs[2], fn_name, "end_snapshot", false)?;
    if start_snapshot > end_snapshot {
        return plan_err!(
            "start_snapshot ({}) must be less than or equal to end_snapshot ({})",
            start_snapshot,
            end_snapshot
        );
    }

    // A change feed reports against the schema as of the window's END snapshot —
    // the table and its schema included, not only its columns. Resolving them at
    // the CURRENT snapshot instead makes a window over a since-dropped table
    // unreadable, even though every snapshot in it still has the table, and makes
    // a window that predates a table's creation resolve to a table it should not
    // see. Official DuckLake resolves at the end snapshot and reports "does not
    // exist at version N" beyond it.
    let (schema_name, table_name_only) = parse_table_name(&table_name);
    let schema = provider
        .get_schema_by_name(schema_name, end_snapshot)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?
        .ok_or_else(|| {
            datafusion::error::DataFusionError::Plan(format!(
                "Schema '{schema_name}' does not exist at snapshot {end_snapshot}"
            ))
        })?;
    let table = provider
        .get_table_by_name(schema.schema_id, table_name_only, end_snapshot)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?
        .ok_or_else(|| {
            datafusion::error::DataFusionError::Plan(format!(
                "Table '{schema_name}.{table_name_only}' does not exist at snapshot \
                 {end_snapshot}"
            ))
        })?;
    let data_path = provider
        .get_data_path()
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
    let (object_store_url, catalog_path) = parse_object_store_url(&data_path)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
    let schema_path = resolve_path(&catalog_path, &schema.path, schema.path_is_relative)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
    let table_path = resolve_path(&schema_path, &table.path, table.path_is_relative)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
    let columns = provider
        .get_table_structure(table.table_id, end_snapshot)
        .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?;
    let table_schema = Arc::new(
        build_arrow_schema(&columns)
            .map_err(|e| datafusion::error::DataFusionError::External(Box::new(e)))?,
    );

    Ok((
        start_snapshot,
        end_snapshot,
        CdcTableContext {
            table_id: table.table_id,
            object_store_url,
            table_path,
            table_schema,
            columns,
        },
    ))
}

impl TableFunctionImpl for DucklakeTableChangesFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        let (start_snapshot, end_snapshot, ctx) =
            parse_cdc_args(&self.provider, exprs, "ducklake_table_changes")?;
        Ok(Arc::new(
            TableChangesTable::new(
                self.provider.clone(),
                ctx.table_id,
                start_snapshot,
                end_snapshot,
                Arc::new(ctx.object_store_url),
                ctx.table_path,
                ctx.table_schema,
            )
            .with_columns(ctx.columns),
        ))
    }
}

#[derive(Debug)]
pub struct DucklakeTableDeletionsFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeTableDeletionsFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeTableDeletionsFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        let (start_snapshot, end_snapshot, ctx) =
            parse_cdc_args(&self.provider, exprs, "ducklake_table_deletions")?;
        Ok(Arc::new(
            TableDeletionsTable::new(
                self.provider.clone(),
                ctx.table_id,
                start_snapshot,
                end_snapshot,
                Arc::new(ctx.object_store_url),
                ctx.table_path,
                ctx.table_schema,
            )
            .with_columns(ctx.columns),
        ))
    }
}

/// `ducklake_table_insertions('schema.table', start, end)`: every row added
/// in the inclusive snapshot window, with `(snapshot_id, rowid)` leading and
/// no `change_type` column — official DuckLake's insertions feed.
#[derive(Debug)]
pub struct DucklakeTableInsertionsFunction {
    provider: Arc<dyn MetadataProvider>,
}

impl DucklakeTableInsertionsFunction {
    pub fn new(provider: Arc<dyn MetadataProvider>) -> Self {
        Self {
            provider,
        }
    }
}

impl TableFunctionImpl for DucklakeTableInsertionsFunction {
    fn call(&self, exprs: &[Expr]) -> DataFusionResult<Arc<dyn TableProvider>> {
        let (start_snapshot, end_snapshot, ctx) =
            parse_cdc_args(&self.provider, exprs, "ducklake_table_insertions")?;
        Ok(Arc::new(
            TableInsertionsTable::new(
                self.provider.clone(),
                ctx.table_id,
                start_snapshot,
                end_snapshot,
                Arc::new(ctx.object_store_url),
                ctx.table_path,
                ctx.table_schema,
            )
            .with_columns(ctx.columns),
        ))
    }
}

/// Registers all ducklake_*() table functions with a SessionContext.
pub fn register_ducklake_functions(
    ctx: &datafusion::execution::context::SessionContext,
    provider: Arc<dyn MetadataProvider>,
) {
    ctx.register_udtf(
        "ducklake_snapshots",
        Arc::new(DucklakeSnapshotsFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_table_info",
        Arc::new(DucklakeTableInfoFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_list_files",
        Arc::new(DucklakeListFilesFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_table_at",
        Arc::new(DucklakeTableAtFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_table_changes",
        Arc::new(DucklakeTableChangesFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_table_insertions",
        Arc::new(DucklakeTableInsertionsFunction::new(provider.clone())),
    );
    ctx.register_udtf(
        "ducklake_table_deletions",
        Arc::new(DucklakeTableDeletionsFunction::new(provider.clone())),
    );
}

#[cfg(test)]
mod snapshot_bound_tests {
    use crate::metadata_provider::parse_snapshot_timestamp;

    #[test]
    fn parses_backend_snapshot_time_formats() {
        for raw in [
            "2026-07-16 12:34:56.123456",
            "2026-07-16 12:34:56",
            "2026-07-16T12:34:56.123456",
            "2026-07-16 12:34:56+00",
            "2026-07-16 12:34:56+00:00",
            "2026-07-16 12:34:56 UTC",
            "2026-07-16T12:34:56Z",
        ] {
            assert!(
                parse_snapshot_timestamp(raw).is_some(),
                "failed to parse {raw:?}"
            );
        }
        // A bare date is midnight UTC.
        assert_eq!(
            parse_snapshot_timestamp("2026-07-16").unwrap(),
            parse_snapshot_timestamp("2026-07-16 00:00:00").unwrap()
        );
    }

    #[test]
    fn rejects_garbage() {
        assert!(parse_snapshot_timestamp("not a time").is_none());
        assert!(parse_snapshot_timestamp("").is_none());
    }
}