buoyant_kernel 0.21.103

//! Defines [`LogReplayScanner`] used by [`TableChangesScan`] to process commit files and extract
//! the metadata needed to generate the Change Data Feed.

use std::collections::{HashMap, HashSet};
use std::slice;
use std::sync::{Arc, LazyLock};

use itertools::Itertools;
use tracing::info;

use crate::actions::visitors::{visit_deletion_vector_at, InCommitTimestampVisitor};
use crate::actions::{
    get_log_add_schema, Add, Cdc, Metadata, Protocol, Remove, ADD_NAME, CDC_NAME, COMMIT_INFO_NAME,
    METADATA_NAME, PROTOCOL_NAME, REMOVE_NAME,
};
use crate::engine_data::{GetData, TypedGetData};
use crate::expressions::{column_expr, column_expr_ref, column_name, ColumnName, Expression};
use crate::path::{AsUrl, ParsedLogPath};
use crate::scan::data_skipping::stats_schema::build_stats_schema;
use crate::scan::data_skipping::DataSkippingFilter;
use crate::scan::state::DvInfo;
use crate::schema::{
    ColumnNamesAndTypes, DataType, SchemaRef, StructField, StructType, ToSchema as _,
};
use crate::table_changes::scan_file::{cdf_scan_row_expression, cdf_scan_row_schema};
use crate::table_configuration::TableConfiguration;
use crate::table_features::{format_features, Operation, TableFeature};
use crate::utils::require;
use crate::{DeltaResult, Engine, EngineData, Error, PredicateRef, RowVisitor};

#[cfg(test)]
mod tests;

/// Scan metadata for a Change Data Feed query. This holds metadata that's needed to read data rows.
pub(crate) struct TableChangesScanMetadata {
    /// Engine data with the schema defined in [`scan_row_schema`]
    ///
    /// Note: The schema of the engine data will be updated in the future to include columns
    /// used by Change Data Feed.
    pub(crate) scan_metadata: Box<dyn EngineData>,
    /// The selection vector used to filter the `scan_metadata`.
    pub(crate) selection_vector: Vec<bool>,
    /// A map from a remove action's path to its deletion vector
    pub(crate) remove_dvs: Arc<HashMap<String, DvInfo>>,
}

/// Given an iterator of [`ParsedLogPath`] returns an iterator of [`TableChangesScanMetadata`].
/// Each row that is selected in the returned `TableChangesScanMetadata.scan_metadata` (according
/// to the `selection_vector` field) _must_ be processed to complete the scan. Non-selected
/// rows _must_ be ignored.
///
/// Note: The [`ParsedLogPath`]s in the `commit_files` iterator must be ordered, contiguous
/// (JSON) commit files.
pub(crate) fn table_changes_action_iter(
    engine: Arc<dyn Engine>,
    start_table_configuration: &TableConfiguration,
    commit_files: impl IntoIterator<Item = ParsedLogPath>,
    table_schema: SchemaRef,
    physical_predicate: Option<(PredicateRef, SchemaRef)>,
) -> DeltaResult<impl Iterator<Item = DeltaResult<TableChangesScanMetadata>>> {
    let filter = physical_predicate
        .and_then(|(predicate, ref_schema)| {
            let stats_schema = build_stats_schema(&ref_schema)?;

            // Parse JSON stats from the raw action batch's `add.stats` column. Unlike the scan
            // path (which transforms first and reads pre-parsed stats), table_changes must
            // resolve deletion vector pairs before filtering, so it operates on raw batches.
            let stats_expr = Arc::new(Expression::parse_json(
                column_expr!("add.stats"),
                stats_schema.clone(),
            ));
            DataSkippingFilter::new(
                engine.as_ref(),
                Some(predicate),
                Some(&stats_schema),
                stats_expr,
                None, // no partition columns for table changes (partition_expr unused)
                column_expr_ref!("partitionValues_parsed"),
                get_log_add_schema().clone(),
                None, // Table changes doesn't use metrics yet
            )
        })
        .map(Arc::new);

    let mut current_configuration = start_table_configuration.clone();
    let result = commit_files
        .into_iter()
        .map(move |commit_file| -> DeltaResult<_> {
            let scanner = LogReplayScanner::try_new(
                engine.as_ref(),
                &mut current_configuration,
                commit_file,
                &table_schema,
            )?;
            scanner.into_scan_batches(engine.clone(), filter.clone())
        }) //Iterator-Result-Iterator-Result
        .flatten_ok() // Iterator-Result-Result
        .map(|x| x?); // Iterator-Result
    Ok(result)
}

/// Processes a single commit file from the log to generate an iterator of
/// [`TableChangesScanMetadata`]. The scanner operates in two phases that _must_ be performed in the
/// following order:
/// 1. Prepare phase [`LogReplayScanner::try_new`]: This iterates over every action in the commit.
///    In this phase, we do the following:
///     - Determine if there exist any `cdc` actions. We determine this in the first phase because
///       the selection vectors for actions are lazily constructed in phase 2. We must know ahead of
///       time whether to filter out add/remove actions.
///     - Constructs the remove deletion vector map from paths belonging to `remove` actions to the
///       action's corresponding [`DvInfo`]. This map will be filtered to only contain paths that
///       exists in another `add` action _within the same commit_. We store the result in
///       `remove_dvs`. Deletion vector resolution affects whether a remove action is selected in
///       the second phase, so we must perform it ahead of time in phase 1.
///     - Ensure that reading is supported on any protocol updates.
///     - Ensure that Change Data Feed is enabled for any metadata update. See  [`TableProperties`]
///     - Ensure that any schema update is compatible with the provided `schema`. Currently, schema
///       compatibility is checked through schema equality. This will be expanded in the future to
///       allow limited schema evolution.
///
/// Note: We check the protocol, change data feed enablement, and schema compatibility in phase 1
/// in order to detect errors and fail early.
///
/// Note: The reader feature [`ReaderFeatures::DeletionVectors`] controls whether the table is
/// allowed to contain deletion vectors. [`TableProperties`].enable_deletion_vectors only
/// determines whether writers are allowed to create _new_ deletion vectors. Hence, we do not need
/// to check the table property for deletion vector enablement.
///
/// See https://github.com/delta-io/delta/blob/master/PROTOCOL.md#deletion-vectors
///
/// TODO: When the kernel supports in-commit timestamps, we will also have to inspect CommitInfo
/// actions to find the timestamp. These are generated when incommit timestamps is enabled.
/// This must be done in the first phase because the second phase lazily transforms engine data with
/// an extra timestamp column. Thus, the timestamp must be known ahead of time.
/// See https://github.com/delta-io/delta-kernel-rs/issues/559
///
/// 2. Scan file generation phase [`LogReplayScanner::into_scan_batches`]: This iterates over every
///    action in the commit, and generates [`TableChangesScanMetadata`]. It does so by transforming
///    the actions using [`add_transform_expr`], and generating selection vectors with the following
///    rules:
///     - If a `cdc` action was found in the prepare phase, only `cdc` actions are selected
///     - Otherwise, select `add` and `remove` actions. Note that only `remove` actions that do not
///       share a path with an `add` action are selected.
///
/// Note: As a consequence of the two phases, LogReplayScanner will iterate over each action in the
/// commit twice. It also may use an unbounded amount of memory, proportional to the number of
/// `add` + `remove` actions in the _single_ commit.
struct LogReplayScanner {
    // True if a `cdc` action was found after running [`LogReplayScanner::try_new`]
    has_cdc_action: bool,
    // A map from path to the deletion vector from the remove action. It is guaranteed that there
    // is an add action with the same path in this commit
    remove_dvs: HashMap<String, DvInfo>,
    // The commit file that this replay scanner will operate on.
    commit_file: ParsedLogPath,
    // The timestamp associated with this commit. This is the file modification time
    // from the commit's [`FileMeta`].
    //
    //
    // TODO when incommit timestamps are supported: If there is a [`CommitInfo`] with a timestamp
    // generated by in-commit timestamps, that timestamp will be used instead.
    //
    // Note: This will be used once an expression is introduced to transform the engine data in
    // [`TableChangesScanMetadata`]
    timestamp: i64,
}

impl LogReplayScanner {
    /// Constructs a LogReplayScanner, performing the Prepare phase detailed in
    /// [`LogReplayScanner`]. This iterates over each action in the commit. It performs the
    /// following:
    /// 1. Check the commits for the presence of a `cdc` action.
    /// 2. Construct a map from path to deletion vector of remove actions that share the same path
    ///    as an add action.
    /// 3. Perform validation on each protocol and metadata action in the commit.
    ///
    /// For more details, see the documentation for [`LogReplayScanner`].
    fn try_new(
        engine: &dyn Engine,
        table_configuration: &mut TableConfiguration,
        commit_file: ParsedLogPath,
        table_schema: &SchemaRef,
    ) -> DeltaResult<Self> {
        let visitor_schema = PreparePhaseVisitor::schema();

        // Note: We do not perform data skipping yet because we need to visit all add and
        // remove actions for deletion vector resolution to be correct.
        //
        // Consider a scenario with a pair of add/remove actions with the same path. The add
        // action has file statistics, while the remove action does not (stats is optional for
        // remove). In this scenario we might skip the add action, while the remove action remains.
        // As a result, we would read the file path for the remove action, which is unnecessary
        // because all of the rows will be filtered by the predicate. Instead, we wait until
        // deletion vectors are resolved so that we can skip both actions in the pair.
        let mut action_iter = engine
            .json_handler()
            .read_json_files(
                slice::from_ref(&commit_file.location),
                visitor_schema,
                None, // not safe to apply data skipping yet
            )?
            .peekable();

        let mut in_commit_timestamp_opt = None;
        if let Some(Ok(actions)) = action_iter.peek() {
            let mut visitor = InCommitTimestampVisitor::default();
            visitor.visit_rows_of(actions.as_ref())?;
            in_commit_timestamp_opt = visitor.in_commit_timestamp;
        }

        let mut remove_dvs = HashMap::default();
        let mut add_paths = HashSet::default();
        let mut has_cdc_action = false;

        for actions in action_iter {
            let actions = actions?;

            let mut visitor = PreparePhaseVisitor {
                add_paths: &mut add_paths,
                remove_dvs: &mut remove_dvs,
                has_cdc_action: &mut has_cdc_action,
            };
            visitor.visit_rows_of(actions.as_ref())?;

            let metadata_opt = Metadata::try_new_from_data(actions.as_ref())?;
            let has_metadata_update = metadata_opt.is_some();
            let protocol_opt = Protocol::try_new_from_data(actions.as_ref())?;
            let has_protocol_update = protocol_opt.is_some();

            if let Some(ref metadata) = metadata_opt {
                let schema = metadata.parse_schema()?;
                // Currently, schema compatibility is defined as having equal schema types. In the
                // future, more permisive schema evolution will be supported.
                // See: https://github.com/delta-io/delta-kernel-rs/issues/523
                require!(
                    table_schema.as_ref() == &schema,
                    Error::change_data_feed_incompatible_schema(table_schema, &schema)
                );
            }

            // Update table configuration with any new Protocol or Metadata from this commit
            if has_metadata_update || has_protocol_update {
                *table_configuration = TableConfiguration::try_new_from(
                    table_configuration,
                    metadata_opt,
                    protocol_opt,
                    commit_file.version,
                )?;

                let writer_features_str = table_configuration
                    .protocol()
                    .writer_features()
                    .map(format_features)
                    .unwrap_or_else(|| "[]".to_string());

                info!(
                    version = commit_file.version,
                    id = table_configuration.metadata().id(),
                    // Writer features are always a superset of reader features, so we log writer features to trace the full set of table features.
                    writerFeatures = %writer_features_str,
                    minReaderVersion = table_configuration.protocol().min_reader_version(),
                    minWriterVersion = table_configuration.protocol().min_writer_version(),
                    schemaString = %table_configuration.metadata().schema_string(),
                    configuration = ?table_configuration.metadata().configuration(),
                    "Table configuration updated during CDF query"
                );
            }

            // If metadata is updated, check if Change Data Feed is enabled
            if has_metadata_update {
                require!(
                    table_configuration.is_feature_enabled(&TableFeature::ChangeDataFeed),
                    Error::change_data_feed_unsupported(commit_file.version)
                );
            }

            // If protocol is updated, check if Change Data Feed is supported
            if has_protocol_update {
                table_configuration
                    .ensure_operation_supported(Operation::Cdf)
                    .map_err(|_| Error::change_data_feed_unsupported(commit_file.version))?;
            }
        }
        // We resolve the remove deletion vector map after visiting the entire commit.
        if has_cdc_action {
            remove_dvs.clear();
        } else {
            // The only (path, deletion_vector) pairs we must track are ones whose path is the
            // same as an `add` action.
            remove_dvs.retain(|rm_path, _| add_paths.contains(rm_path));
        }

        // If ICT is enabled, then set the timestamp to be the ICT; otherwise, default to the
        // last_modified timestamp value
        let timestamp = if table_configuration.is_feature_enabled(&TableFeature::InCommitTimestamp)
        {
            let Some(in_commit_timestamp) = in_commit_timestamp_opt else {
                return Err(Error::generic(format!(
                    "In-commit timestamp is enabled but not found in commit at version {}",
                    commit_file.version
                )));
            };
            in_commit_timestamp
        } else {
            commit_file.location.last_modified
        };

        info!(
            version = commit_file.version,
            id = table_configuration.metadata().id(),
            remove_dvs_size = remove_dvs.len(),
            has_cdc_action = has_cdc_action,
            file_path = %commit_file.location.as_url(),
            timestamp = timestamp,
            "Phase 1 of CDF query processing completed"
        );

        Ok(LogReplayScanner {
            timestamp,
            commit_file,
            has_cdc_action,
            remove_dvs,
        })
    }
    /// Generates an iterator of [`TableChangesScanMetadata`] by iterating over each action of the
    /// commit, generating a selection vector, and transforming the engine data. This performs
    /// phase 2 of [`LogReplayScanner`].
    fn into_scan_batches(
        self,
        engine: Arc<dyn Engine>,
        filter: Option<Arc<DataSkippingFilter>>,
    ) -> DeltaResult<impl Iterator<Item = DeltaResult<TableChangesScanMetadata>>> {
        let Self {
            has_cdc_action,
            remove_dvs,
            commit_file,
            // TODO: Add the timestamp as a column with an expression
            timestamp,
        } = self;
        let remove_dvs = Arc::new(remove_dvs);

        let schema = FileActionSelectionVisitor::schema();
        let action_iter = engine.json_handler().read_json_files(
            slice::from_ref(&commit_file.location),
            schema,
            None,
        )?;
        let commit_version = commit_file
            .version
            .try_into()
            .map_err(|_| Error::generic("Failed to convert commit version to i64"))?;
        let evaluator = engine.evaluation_handler().new_expression_evaluator(
            get_log_add_schema().clone(),
            Arc::new(cdf_scan_row_expression(timestamp, commit_version)),
            cdf_scan_row_schema().into(),
        )?;

        let result = action_iter.map(move |actions| -> DeltaResult<_> {
            let actions = actions?;

            // Apply data skipping to get back a selection vector for actions that passed skipping.
            // We start our selection vector based on what was filtered. We will add to this vector
            // below if a file has been removed. Note: None implies all files passed data skipping.
            let selection_vector = match &filter {
                Some(filter) => filter.apply(actions.as_ref())?,
                None => vec![true; actions.len()],
            };

            let mut visitor =
                FileActionSelectionVisitor::new(&remove_dvs, selection_vector, has_cdc_action);
            visitor.visit_rows_of(actions.as_ref())?;
            let scan_metadata = evaluator.evaluate(actions.as_ref())?;
            Ok(TableChangesScanMetadata {
                scan_metadata,
                selection_vector: visitor.selection_vector,
                remove_dvs: remove_dvs.clone(),
            })
        });
        Ok(result)
    }
}

// This is a visitor used in the prepare phase of [`LogReplayScanner`]. See
// [`LogReplayScanner::try_new`] for details usage.
struct PreparePhaseVisitor<'a> {
    has_cdc_action: &'a mut bool,
    add_paths: &'a mut HashSet<String>,
    remove_dvs: &'a mut HashMap<String, DvInfo>,
}
impl PreparePhaseVisitor<'_> {
    fn schema() -> Arc<StructType> {
        Arc::new(StructType::new_unchecked(vec![
            StructField::nullable(ADD_NAME, Add::to_schema()),
            StructField::nullable(REMOVE_NAME, Remove::to_schema()),
            StructField::nullable(CDC_NAME, Cdc::to_schema()),
            StructField::nullable(METADATA_NAME, Metadata::to_schema()),
            StructField::nullable(PROTOCOL_NAME, Protocol::to_schema()),
            StructField::nullable(
                COMMIT_INFO_NAME,
                StructType::new_unchecked([StructField::new(
                    "inCommitTimestamp",
                    DataType::LONG,
                    true,
                )]),
            ),
        ]))
    }
}

impl RowVisitor for PreparePhaseVisitor<'_> {
    fn selected_column_names_and_types(&self) -> (&'static [ColumnName], &'static [DataType]) {
        // NOTE: The order of the names and types is based on [`PreparePhaseVisitor::schema`]
        static NAMES_AND_TYPES: LazyLock<ColumnNamesAndTypes> = LazyLock::new(|| {
            const STRING: DataType = DataType::STRING;
            const INTEGER: DataType = DataType::INTEGER;
            const LONG: DataType = DataType::LONG;
            const BOOLEAN: DataType = DataType::BOOLEAN;
            let types_and_names = vec![
                (STRING, column_name!("add.path")),
                (BOOLEAN, column_name!("add.dataChange")),
                (STRING, column_name!("remove.path")),
                (BOOLEAN, column_name!("remove.dataChange")),
                (STRING, column_name!("remove.deletionVector.storageType")),
                (STRING, column_name!("remove.deletionVector.pathOrInlineDv")),
                (INTEGER, column_name!("remove.deletionVector.offset")),
                (INTEGER, column_name!("remove.deletionVector.sizeInBytes")),
                (LONG, column_name!("remove.deletionVector.cardinality")),
                (STRING, column_name!("cdc.path")),
                (LONG, column_name!("commitInfo.inCommitTimestamp")),
            ];
            let (types, names) = types_and_names.into_iter().unzip();
            (names, types).into()
        });
        NAMES_AND_TYPES.as_ref()
    }

    fn visit<'b>(&mut self, row_count: usize, getters: &[&'b dyn GetData<'b>]) -> DeltaResult<()> {
        require!(
            getters.len() == 11,
            Error::InternalError(format!(
                "Wrong number of PreparePhaseVisitor getters: {}",
                getters.len()
            ))
        );
        for i in 0..row_count {
            if let Some(path) = getters[0].get_str(i, "add.path")? {
                // If no data was changed, we must ignore that action
                if !*self.has_cdc_action && getters[1].get(i, "add.dataChange")? {
                    self.add_paths.insert(path.to_string());
                }
            } else if let Some(path) = getters[2].get_str(i, "remove.path")? {
                // If no data was changed, we must ignore that action
                if !*self.has_cdc_action && getters[3].get(i, "remove.dataChange")? {
                    let deletion_vector = visit_deletion_vector_at(i, &getters[4..=8])?;
                    self.remove_dvs
                        .insert(path.to_string(), DvInfo { deletion_vector });
                }
            } else if getters[9].get_str(i, "cdc.path")?.is_some() {
                *self.has_cdc_action = true;
            }
        }
        Ok(())
    }
}

// This visitor generates selection vectors based on the rules specified in [`LogReplayScanner`].
// See [`LogReplayScanner::into_scan_batches`] for usage.
struct FileActionSelectionVisitor<'a> {
    selection_vector: Vec<bool>,
    has_cdc_action: bool,
    remove_dvs: &'a HashMap<String, DvInfo>,
}

impl<'a> FileActionSelectionVisitor<'a> {
    fn new(
        remove_dvs: &'a HashMap<String, DvInfo>,
        selection_vector: Vec<bool>,
        has_cdc_action: bool,
    ) -> Self {
        FileActionSelectionVisitor {
            selection_vector,
            has_cdc_action,
            remove_dvs,
        }
    }
    fn schema() -> Arc<StructType> {
        Arc::new(StructType::new_unchecked(vec![
            StructField::nullable(CDC_NAME, Cdc::to_schema()),
            StructField::nullable(ADD_NAME, Add::to_schema()),
            StructField::nullable(REMOVE_NAME, Remove::to_schema()),
        ]))
    }
}

impl RowVisitor for FileActionSelectionVisitor<'_> {
    fn selected_column_names_and_types(&self) -> (&'static [ColumnName], &'static [DataType]) {
        // Note: The order of the names and types is based on [`FileActionSelectionVisitor::schema`]
        static NAMES_AND_TYPES: LazyLock<ColumnNamesAndTypes> = LazyLock::new(|| {
            const STRING: DataType = DataType::STRING;
            const BOOLEAN: DataType = DataType::BOOLEAN;
            let types_and_names = vec![
                (STRING, column_name!("cdc.path")),
                (STRING, column_name!("add.path")),
                (BOOLEAN, column_name!("add.dataChange")),
                (STRING, column_name!("remove.path")),
                (BOOLEAN, column_name!("remove.dataChange")),
            ];
            let (types, names) = types_and_names.into_iter().unzip();
            (names, types).into()
        });
        NAMES_AND_TYPES.as_ref()
    }

    fn visit<'b>(&mut self, row_count: usize, getters: &[&'b dyn GetData<'b>]) -> DeltaResult<()> {
        require!(
            getters.len() == 5,
            Error::InternalError(format!(
                "Wrong number of FileActionSelectionVisitor getters: {}",
                getters.len()
            ))
        );

        for i in 0..row_count {
            if !self.selection_vector[i] {
                continue;
            }

            if self.has_cdc_action {
                self.selection_vector[i] = getters[0].get_str(i, "cdc.path")?.is_some()
            } else if getters[1].get_str(i, "add.path")?.is_some() {
                self.selection_vector[i] = getters[2].get(i, "add.dataChange")?;
            } else if let Some(path) = getters[3].get_str(i, "remove.path")? {
                let data_change: bool = getters[4].get(i, "remove.dataChange")?;
                self.selection_vector[i] = data_change && !self.remove_dvs.contains_key(path)
            } else {
                self.selection_vector[i] = false
            }
        }
        Ok(())
    }
}