delta_kernel 0.30.1

// Allow unreachable_pub because this module is pub when test-utils is enabled,
// making the pub items reachable from integration tests.
#![allow(unreachable_pub)]

use std::sync::LazyLock;

use serde::{Deserialize, Serialize};

use crate::actions::{DomainMetadata, NUM_RECORDS};
use crate::engine_data::{GetData, RowVisitor, TypedGetData as _};
use crate::schema::{column_name, ColumnName, ColumnNamesAndTypes, DataType};
use crate::utils::require;
use crate::{KernelError, KernelResult};

#[derive(Debug, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RowTrackingDomainMetadata {
    // NB: The Delta spec does not rule out negative high water marks
    row_id_high_water_mark: i64,
}

/// The domain name for row tracking metadata.
pub(crate) const ROW_TRACKING_DOMAIN_NAME: &str = "delta.rowTracking";

/// The row-tracking high-water mark before any fresh row IDs have been assigned.
pub const ROW_TRACKING_INITIAL_HIGH_WATER_MARK: i64 = -1;

impl RowTrackingDomainMetadata {
    pub(crate) fn new(row_id_high_water_mark: i64) -> Self {
        RowTrackingDomainMetadata {
            row_id_high_water_mark,
        }
    }

    /// Returns the highest row ID represented by this metadata.
    pub(crate) fn high_water_mark(&self) -> i64 {
        self.row_id_high_water_mark
    }

    /// Creates the initial row tracking domain metadata for a newly created table.
    ///
    /// Sets the high-water mark to [`ROW_TRACKING_INITIAL_HIGH_WATER_MARK`], meaning no rows have
    /// been assigned IDs yet. The first file written will receive `baseRowId = 0`.
    pub(crate) fn initial() -> Self {
        Self::new(ROW_TRACKING_INITIAL_HIGH_WATER_MARK)
    }
}

pub(crate) fn parse_row_tracking_high_water_mark(configuration: &str) -> KernelResult<i64> {
    Ok(serde_json::from_str::<RowTrackingDomainMetadata>(configuration)?.high_water_mark())
}

impl TryFrom<RowTrackingDomainMetadata> for DomainMetadata {
    type Error = crate::KernelError;

    fn try_from(metadata: RowTrackingDomainMetadata) -> KernelResult<Self> {
        Ok(DomainMetadata::new(
            ROW_TRACKING_DOMAIN_NAME.to_string(),
            serde_json::to_string(&metadata)?,
        ))
    }
}

/// A row visitor that iterates over preliminary [`Add`] actions as returned by the engine and
/// computes a base row ID for each action.
/// It expects to visit engine data with a nested field 'stats.numRecords' which is
/// part of a Delta add action.
///
/// This visitor is only required for the row tracking write path. The read path will be completely
/// implemented via expressions.
///
/// [`Add`]: delta_kernel::actions::Add
pub(crate) struct RowTrackingVisitor {
    /// High water mark for row IDs
    pub(crate) row_id_high_water_mark: i64,

    collect_base_row_ids: bool,
    base_row_ids: Vec<i64>,
}

impl RowTrackingVisitor {
    pub(crate) fn new(row_id_high_water_mark: Option<i64>) -> Self {
        // A table might not have a row ID high water mark yet, so we model the input as an
        // Option<i64>
        Self {
            row_id_high_water_mark: row_id_high_water_mark
                .unwrap_or(ROW_TRACKING_INITIAL_HIGH_WATER_MARK),
            collect_base_row_ids: true,
            base_row_ids: Vec::new(),
        }
    }

    /// Like [`Self::new`], but does not materialize rowid bases.
    pub(crate) fn dry_run(row_id_high_water_mark: Option<i64>) -> Self {
        Self {
            collect_base_row_ids: false,
            ..Self::new(row_id_high_water_mark)
        }
    }

    /// Takes all rowid bases materialized so far.
    pub(crate) fn take_base_row_ids(&mut self) -> Vec<i64> {
        std::mem::take(&mut self.base_row_ids)
    }
}

impl RowVisitor for RowTrackingVisitor {
    fn selected_column_names_and_types(&self) -> (&'static [ColumnName], &'static [DataType]) {
        static NAMES_AND_TYPES: LazyLock<ColumnNamesAndTypes> = LazyLock::new(|| {
            (
                vec![column_name!("stats", NUM_RECORDS)],
                vec![DataType::LONG],
            )
                .into()
        });
        NAMES_AND_TYPES.as_ref()
    }

    fn visit<'a>(&mut self, row_count: usize, getters: &[&'a dyn GetData<'a>]) -> KernelResult<()> {
        require!(
            getters.len() == 1,
            KernelError::generic(format!(
                "Wrong number of RowTrackingVisitor getters: {}",
                getters.len()
            ))
        );

        if self.collect_base_row_ids {
            self.base_row_ids.reserve(row_count);
        }

        for i in 0..row_count {
            let num_records: i64 = getters[0].get_opt(i, NUM_RECORDS)?.ok_or_else(|| {
                KernelError::InternalError(format!(
                    "{NUM_RECORDS} must be present in Add actions when row tracking is enabled."
                ))
            })?;
            if self.collect_base_row_ids {
                self.base_row_ids.push(self.row_id_high_water_mark + 1);
            }
            self.row_id_high_water_mark += num_records;
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use derive_more::Constructor;

    use super::*;
    use crate::engine_data::GetData;
    use crate::unit_test_utils::assert_result_error_with_message;

    /// Mock GetData implementation for testing
    #[derive(Constructor)]
    struct MockGetData {
        num_records_values: Vec<Option<i64>>,
    }

    impl<'a> GetData<'a> for MockGetData {
        fn get_long(&'a self, row_index: usize, field_name: &str) -> KernelResult<Option<i64>> {
            if field_name == NUM_RECORDS {
                Ok(self.num_records_values.get(row_index).copied().flatten())
            } else {
                Ok(None)
            }
        }
    }

    fn create_getters(num_records_mock: &MockGetData) -> Vec<&dyn GetData<'_>> {
        vec![num_records_mock]
    }

    #[test]
    fn test_visit_basic_functionality() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(None);
        let num_records_mock = MockGetData::new(vec![Some(10), Some(5), Some(20)]);
        let getters = create_getters(&num_records_mock);

        visitor.visit(3, &getters)?;

        // Check that base row IDs are calculated correctly
        assert_eq!(visitor.take_base_row_ids(), vec![0, 10, 15]);

        // Check that high water mark is updated correctly
        assert_eq!(visitor.row_id_high_water_mark, 34); // -1 + 10 + 5 + 20

        Ok(())
    }

    #[test]
    fn test_visit_with_negative_high_water_mark() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(-5));
        let num_records_mock = MockGetData::new(vec![Some(3), Some(2)]);
        let getters = create_getters(&num_records_mock);

        visitor.visit(2, &getters)?;

        // Base row IDs should start from high_water_mark + 1
        assert_eq!(visitor.take_base_row_ids(), vec![-4, -1]);

        // High water mark should be updated
        assert_eq!(visitor.row_id_high_water_mark, 0); // -5 + 3 + 2 = 0

        Ok(())
    }

    #[test]
    fn test_visit_with_zero_records() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(10));
        let num_records_mock = MockGetData::new(vec![Some(0), Some(0), Some(5)]);
        let getters = create_getters(&num_records_mock);

        visitor.visit(3, &getters)?;

        // Base row IDs should still be assigned even for zero-record files
        assert_eq!(visitor.take_base_row_ids(), vec![11, 11, 11]);

        // High water mark should only increase by non-zero records
        assert_eq!(visitor.row_id_high_water_mark, 15); // 10 + 0 + 0 + 5

        Ok(())
    }

    #[test]
    fn test_visit_empty_batch() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(42));
        let num_records_mock = MockGetData::new(vec![]);
        let getters = create_getters(&num_records_mock);

        visitor.visit(0, &getters)?;

        // Should handle empty batch gracefully
        assert!(visitor.take_base_row_ids().is_empty());
        assert_eq!(visitor.row_id_high_water_mark, 42); // Should remain unchanged

        Ok(())
    }

    #[test]
    fn test_visit_multiple_batches() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(0));

        // First batch
        let num_records_mock1 = MockGetData::new(vec![Some(10), Some(5)]);
        let getters1 = create_getters(&num_records_mock1);
        visitor.visit(2, &getters1)?;
        let first_base_row_ids = visitor.take_base_row_ids();

        // Second batch
        let num_records_mock2 = MockGetData::new(vec![Some(3), Some(7), Some(2)]);
        let getters2 = create_getters(&num_records_mock2);
        visitor.visit(3, &getters2)?;
        let second_base_row_ids = visitor.take_base_row_ids();

        // Check first batch: starts at 1, then 11
        assert_eq!(first_base_row_ids, vec![1, 11]);

        // Check second batch: starts at 16, then 19, then 26
        assert_eq!(second_base_row_ids, vec![16, 19, 26]);

        // Check final high water mark: 0 + 10 + 5 + 3 + 7 + 2 = 27
        assert_eq!(visitor.row_id_high_water_mark, 27);

        Ok(())
    }

    #[test]
    fn test_visit_wrong_getter_count() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(0));
        let wrong_getters: Vec<&dyn GetData<'_>> = vec![]; // No getters instead of expected count

        let result = visitor.visit(1, &wrong_getters);
        assert_result_error_with_message(result, "Wrong number of RowTrackingVisitor getters");

        Ok(())
    }

    #[test]
    fn test_visit_missing_num_records() -> KernelResult<()> {
        let mut visitor = RowTrackingVisitor::new(Some(0));
        let num_records_mock = MockGetData::new(vec![None]); // Missing numRecords
        let getters = create_getters(&num_records_mock);

        let result = visitor.visit(1, &getters);
        assert_result_error_with_message(
            result,
            &format!("{NUM_RECORDS} must be present in Add actions when row tracking is enabled"),
        );

        Ok(())
    }

    #[test]
    fn test_selected_column_names_and_types() {
        let visitor = RowTrackingVisitor::new(Some(0));
        let (names, types) = visitor.selected_column_names_and_types();

        assert_eq!(names, (vec![column_name!("stats", NUM_RECORDS)]));
        assert_eq!(types, vec![DataType::LONG]);
    }

    #[test]
    fn test_serialization_roundtrip() -> KernelResult<()> {
        let original = RowTrackingDomainMetadata::new(-42);
        let json = serde_json::to_string(&original)?;
        let deserialized: RowTrackingDomainMetadata = serde_json::from_str(&json)?;

        assert_eq!(
            original.row_id_high_water_mark,
            deserialized.row_id_high_water_mark
        );

        Ok(())
    }
}