delta_kernel 0.25.0

//! Transforms for populating checkpoint-specific fields in the Add action.
//!
//! This module ensures that Add actions in checkpoints have the correct format for:
//!
//! **Statistics** (controlled by `delta.checkpoint.writeStatsAsStruct` / `writeStatsAsJson`):
//! - `stats`: JSON string format (default: enabled)
//! - `stats_parsed`: Native struct format (default: disabled)
//!
//! **Partition values** (controlled by `delta.checkpoint.writeStatsAsStruct`):
//! - `partitionValues`: String-valued map (always present)
//! - `partitionValues_parsed`: Native typed struct (only when `writeStatsAsStruct=true`)
//!
//! This module provides transforms to populate these fields using COALESCE expressions,
//! ensuring that values are preserved regardless of the source format (commits vs checkpoints).

use std::sync::{Arc, LazyLock};

use crate::actions::{ADD_NAME, STATS_PARSED as STATS_PARSED_FIELD};
use crate::expressions::{col, Expression, ExpressionRef, UnaryExpressionOp};
use crate::schema::{DataType, SchemaRef, SchemaStructPatchBuilder, StructField, StructType};
use crate::struct_patch::ProjectionStructPatchBuilder;
use crate::table_properties::TableProperties;
use crate::{DeltaResult, Error};

pub(crate) const STATS_FIELD: &str = "stats";
pub(crate) const PARTITION_VALUES_FIELD: &str = "partitionValues";
pub(crate) const PARTITION_VALUES_PARSED_FIELD: &str = "partitionValues_parsed";

/// Configuration for stats transformation based on table properties.
#[derive(Debug, Clone, Copy)]
pub(crate) struct StatsTransformConfig {
    pub write_stats_as_json: bool,
    pub write_stats_as_struct: bool,
}

impl StatsTransformConfig {
    pub(super) fn from_table_properties(properties: &TableProperties) -> Self {
        Self {
            write_stats_as_json: properties.should_write_stats_as_json(),
            write_stats_as_struct: properties.should_write_stats_as_struct(),
        }
    }
}

/// Builds a transform for the Add action to populate and/or drop stats and partition fields.
///
/// The transform handles statistics based on table properties:
/// - When `writeStatsAsJson=true`: `stats = COALESCE(stats, ToJson(stats_parsed))`
/// - When `writeStatsAsJson=false`: drop `stats` field
/// - When `writeStatsAsStruct=true`: `stats_parsed = COALESCE(stats_parsed, ParseJson(stats))`
/// - When `writeStatsAsStruct=false`: drop `stats_parsed` field
///
/// For partitioned tables when `writeStatsAsStruct=true`, it also populates:
/// - `partitionValues_parsed = COALESCE(partitionValues_parsed, MAP_TO_STRUCT(partitionValues))`
///
/// Returns a top-level transform that wraps the nested Add transform, ensuring the
/// full checkpoint batch is produced with the modified Add action.
///
/// # Arguments
///
/// * `stats_schema` - The expected schema for parsed file statistics, typically generated by
///   [`expected_stats_schema`]. This schema has the following structure:
///
///   ```text
///   {
///     numRecords: long,
///     nullCount: <nested struct with all leaf fields as long>,
///     minValues: <nested struct matching eligible column types>,
///     maxValues: <nested struct matching eligible column types>,
///   }
///   ```
///
///   The schema is derived from the table's physical file schema and table properties
///   (`dataSkippingNumIndexedCols`, `dataSkippingStatsColumns`). Only columns eligible for data
///   skipping are included in `minValues`/`maxValues`.
///
/// * `partition_schema` - The physical partition schema for `partitionValues_parsed`. `None` for
///   non-partitioned tables.
///
/// [`expected_stats_schema`]: crate::scan::data_skipping::stats_schema::expected_stats_schema
pub(crate) fn build_checkpoint_transform(
    config: &StatsTransformConfig,
    read_schema: &StructType,
    stats_schema: &SchemaRef,
    partition_schema: Option<&SchemaRef>,
) -> DeltaResult<(SchemaRef, ExpressionRef)> {
    let mut patch_builder = ProjectionStructPatchBuilder::new(read_schema);

    // Handle stats field
    if config.write_stats_as_json {
        // Populate stats from stats_parsed if needed (for old checkpoints that only had
        // stats_parsed)
        patch_builder =
            patch_builder.replace_expr_at([ADD_NAME], STATS_FIELD, STATS_JSON_EXPR.clone());
    } else {
        // Drop stats field when not writing as JSON
        patch_builder = patch_builder.drop_at([ADD_NAME], STATS_FIELD);
    }

    // Handle stats_parsed field
    // Note: stats_parsed was added to read schema (via build_checkpoint_read_schema),
    // so we always need to either replace it (with COALESCE) or drop it.
    if config.write_stats_as_struct {
        // Populate stats_parsed from JSON stats (for commits that only have JSON stats)
        let stats_parsed_expr = build_stats_parsed_expr(stats_schema);
        patch_builder =
            patch_builder.replace_expr_at([ADD_NAME], STATS_PARSED_FIELD, stats_parsed_expr);
    } else {
        // Drop stats_parsed field when not writing as struct
        patch_builder = patch_builder.drop_at([ADD_NAME], STATS_PARSED_FIELD);
    }

    // Handle partitionValues_parsed field (only for partitioned tables)
    if partition_schema.is_some() {
        if config.write_stats_as_struct {
            let pv_parsed_expr = build_partition_values_parsed_expr();
            patch_builder = patch_builder.replace_expr_at(
                [ADD_NAME],
                PARTITION_VALUES_PARSED_FIELD,
                pv_parsed_expr,
            );
        } else {
            // Drop partitionValues_parsed since it was added to read schema
            patch_builder = patch_builder.drop_at([ADD_NAME], PARTITION_VALUES_PARSED_FIELD);
        }
    }

    patch_builder.build()
}

/// Builds a read schema that includes `stats_parsed` and optionally `partitionValues_parsed`
/// in the Add action.
///
/// The read schema must be union-compatible across all log segment files (checkpoints and
/// JSON commits). This means all reads use the same schema even though commits don't have
/// `stats_parsed` or `partitionValues_parsed` — those columns are read as nulls. This
/// union-compatible schema ensures log replay can process checkpoint and commit batches
/// uniformly, and COALESCE expressions can operate correctly across both sources.
///
/// # Errors
///
/// Returns an error if:
/// - The `add` field is not found or is not a struct type
/// - The `stats_parsed` or `partitionValues_parsed` field already exists in the Add schema
pub(crate) fn build_checkpoint_read_schema(
    base_schema: &StructType,
    stats_schema: &StructType,
    partition_schema: Option<&StructType>,
) -> DeltaResult<SchemaRef> {
    transform_add_schema(base_schema, |add_struct| {
        // Validate fields aren't already present
        if add_struct.field(STATS_PARSED_FIELD).is_some() {
            return Err(Error::generic(
                "stats_parsed field already exists in Add schema",
            ));
        }
        if partition_schema.is_some() && add_struct.field(PARTITION_VALUES_PARSED_FIELD).is_some() {
            return Err(Error::generic(
                "partitionValues_parsed field already exists in Add schema",
            ));
        }
        let mut patch = SchemaStructPatchBuilder::new().insert_after(
            STATS_FIELD,
            StructField::nullable(STATS_PARSED_FIELD, stats_schema.clone()),
        );
        if let Some(pv_schema) = partition_schema {
            patch = patch.insert_after(
                PARTITION_VALUES_FIELD,
                StructField::nullable(PARTITION_VALUES_PARSED_FIELD, pv_schema.clone()),
            );
        }
        patch.build(add_struct)
    })
}

// ========================
// Private helpers
// ========================

/// Builds expression: `stats_parsed = COALESCE(stats_parsed, ParseJson(stats, schema))`
///
/// This expression prefers existing stats_parsed, falling back to parsing JSON stats.
/// If `stats_parsed` is non-null, the data originated from a checkpoint (commits only
/// contain JSON stats, so `stats_parsed` will be null for commit-sourced rows).
///
/// Column paths are relative to the full batch, not the nested Add struct.
fn build_stats_parsed_expr(stats_schema: &SchemaRef) -> ExpressionRef {
    Arc::new(Expression::coalesce([
        col!(ADD_NAME, STATS_PARSED_FIELD),
        Expression::parse_json(col!(ADD_NAME, STATS_FIELD), stats_schema.clone()),
    ]))
}

/// Builds expression: `partitionValues_parsed = COALESCE(partitionValues_parsed,
///     MAP_TO_STRUCT(partitionValues))`
///
/// This expression prefers existing `partitionValues_parsed`, falling back to converting
/// the string-valued `partitionValues` map into a native typed struct. The target struct
/// type (field names and data types) is determined by the output schema — `MAP_TO_STRUCT`
/// itself carries no schema, so the expression evaluator uses the expected output type to
/// parse each string value into the correct native type.
///
/// Column paths are relative to the full batch, not the nested Add struct.
fn build_partition_values_parsed_expr() -> ExpressionRef {
    Arc::new(Expression::coalesce([
        col!(ADD_NAME, PARTITION_VALUES_PARSED_FIELD),
        Expression::map_to_struct(col!(ADD_NAME, PARTITION_VALUES_FIELD)),
    ]))
}

/// Static expression: `stats = COALESCE(stats, ToJson(stats_parsed))`
///
/// This expression prefers existing JSON stats, falling back to converting stats_parsed.
/// Column paths are relative to the full batch, not the nested Add struct.
static STATS_JSON_EXPR: LazyLock<ExpressionRef> = LazyLock::new(|| {
    Arc::new(Expression::coalesce([
        col!(ADD_NAME, STATS_FIELD),
        Expression::unary(
            UnaryExpressionOp::ToJson,
            col!(ADD_NAME, STATS_PARSED_FIELD),
        ),
    ]))
});

/// Transforms the Add action schema within a checkpoint schema.
///
/// This helper applies a transformation function to the Add struct and returns
/// a new schema with the modified Add field.
///
/// # Errors
///
/// Returns an error if:
/// - The `add` field is not found in the schema
/// - The `add` field is not a struct type
fn transform_add_schema(
    base_schema: &StructType,
    transform_fn: impl FnOnce(&StructType) -> DeltaResult<StructType>,
) -> DeltaResult<SchemaRef> {
    // Find and validate the add field
    let add_field = base_schema
        .field(ADD_NAME)
        .ok_or_else(|| Error::generic("Expected 'add' field in checkpoint schema"))?;

    let DataType::Struct(add_struct) = &add_field.data_type else {
        return Err(Error::generic(format!(
            "Expected 'add' field to be a struct type, got {:?}",
            add_field.data_type
        )));
    };

    let modified_add = transform_fn(add_struct)?;
    let new_add_field = StructField::new(ADD_NAME, modified_add, add_field.nullable)
        .with_metadata(add_field.metadata.clone());
    let new_schema = SchemaStructPatchBuilder::new()
        .replace(ADD_NAME, new_add_field)
        .build(base_schema)?;

    Ok(Arc::new(new_schema))
}

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

    use super::*;
    use crate::actions::NUM_RECORDS;
    use crate::expressions::ExpressionStructPatch;
    use crate::schema::MapType;

    #[test]
    fn test_config_defaults() {
        // Default: writeStatsAsJson=true, writeStatsAsStruct=false (per protocol)
        let props = TableProperties::default();
        let config = StatsTransformConfig::from_table_properties(&props);
        assert!(config.write_stats_as_json);
        assert!(!config.write_stats_as_struct);
    }

    #[test]
    fn test_config_with_struct_enabled() {
        let props = TableProperties {
            checkpoint_write_stats_as_struct: Some(true),
            ..Default::default()
        };
        let config = StatsTransformConfig::from_table_properties(&props);
        assert!(config.write_stats_as_json);
        assert!(config.write_stats_as_struct);
    }

    /// Helper to extract the outer and inner patches from a stats transform expression.
    /// Returns (outer_patch, inner_patch).
    fn extract_patches(expr: &Expression) -> (&ExpressionStructPatch, &ExpressionStructPatch) {
        let Expression::StructPatch(outer) = expr else {
            panic!("Expected outer StructPatch expression");
        };

        // Outer should be top-level (no input path)
        assert!(
            outer.input_path.is_none(),
            "Outer patch should be top-level"
        );

        // Outer should replace "add" field
        let add_field_patch = outer
            .field_patches
            .get(ADD_NAME)
            .expect("Outer patch should have 'add' field patch");
        assert!(
            !add_field_patch.keep_input && !add_field_patch.insertions.is_empty(),
            "Should replace 'add' field"
        );

        // Extract inner patch
        let Expression::StructPatch(inner) = add_field_patch.insertions[0].as_ref() else {
            panic!("Expected inner StructPatch expression for 'add' field");
        };

        // Inner should target "add" path
        assert_eq!(
            inner.input_path.as_ref().map(|p| p.to_string()),
            Some("add".to_string()),
            "Inner patch should target 'add' path"
        );

        (outer, inner)
    }

    /// Helper to check if a field patch is a drop (replace with nothing).
    fn is_drop(patch: &ExpressionStructPatch, field: &str) -> bool {
        patch
            .field_patches
            .get(field)
            .is_some_and(|ft| !ft.keep_input && ft.insertions.is_empty())
    }

    /// Helper to check if a field patch is a replacement with an expression.
    fn is_replacement(patch: &ExpressionStructPatch, field: &str) -> bool {
        patch
            .field_patches
            .get(field)
            .is_some_and(|ft| !ft.keep_input && !ft.insertions.is_empty())
    }

    fn build_checkpoint_transform(
        config: &StatsTransformConfig,
        stats_schema: &SchemaRef,
        partition_schema: Option<&SchemaRef>,
    ) -> (SchemaRef, ExpressionRef) {
        let base_schema = StructType::new_unchecked([StructField::nullable(
            ADD_NAME,
            add_schema(partition_schema.is_some()),
        )]);
        let read_schema = build_checkpoint_read_schema(
            &base_schema,
            stats_schema.as_ref(),
            partition_schema.map(AsRef::as_ref),
        )
        .expect("build checkpoint read schema should produce a valid schema");
        super::build_checkpoint_transform(config, &read_schema, stats_schema, partition_schema)
            .expect("build checkpoint transform should produce a valid schema and expression")
    }

    fn add_schema(with_partition_schema: bool) -> StructType {
        let fields = [
            Some(StructField::not_null("path", DataType::STRING)),
            with_partition_schema.then(|| {
                let partition_values = MapType::new(DataType::STRING, DataType::STRING, true);
                StructField::nullable(PARTITION_VALUES_FIELD, partition_values)
            }),
            Some(StructField::nullable(STATS_FIELD, DataType::STRING)),
        ];
        StructType::new_unchecked(fields.into_iter().flatten())
    }

    #[test]
    fn test_build_transform_with_json_only() {
        // writeStatsAsJson=true, writeStatsAsStruct=false (default)
        // Inner patch: stats=COALESCE, stats_parsed=drop
        let config = StatsTransformConfig {
            write_stats_as_json: true,
            write_stats_as_struct: false,
        };
        let stats_schema = Arc::new(StructType::new_unchecked([]));
        let (_, transform_expr) = build_checkpoint_transform(&config, &stats_schema, None);

        let (_, inner) = extract_patches(&transform_expr);

        // stats should be replaced with COALESCE expression
        assert!(
            is_replacement(inner, STATS_FIELD),
            "stats should be replaced"
        );

        // stats_parsed should be dropped
        assert!(
            is_drop(inner, STATS_PARSED_FIELD),
            "stats_parsed should be dropped"
        );
    }

    #[test]
    fn test_build_transform_populates_struct_stats_and_partition_values() {
        let config = StatsTransformConfig {
            write_stats_as_json: true,
            write_stats_as_struct: true,
        };
        let stats_schema = Arc::new(StructType::new_unchecked([]));
        let pv_schema = Arc::new(StructType::new_unchecked([
            StructField::nullable("year", DataType::INTEGER),
            StructField::nullable("month", DataType::INTEGER),
        ]));
        let (_, transform_expr) =
            build_checkpoint_transform(&config, &stats_schema, Some(&pv_schema));

        let (_, inner) = extract_patches(&transform_expr);

        assert!(
            is_replacement(inner, STATS_FIELD),
            "stats should be replaced"
        );
        assert!(
            is_replacement(inner, STATS_PARSED_FIELD),
            "stats_parsed should be replaced"
        );
        assert!(
            is_replacement(inner, PARTITION_VALUES_PARSED_FIELD),
            "partitionValues_parsed should be replaced"
        );
    }

    #[rstest]
    #[case::no_stats(false, false, false, &["path"])]
    #[case::json_only(true, false, false, &["path", "stats"])]
    #[case::struct_only(false, true, false, &["path", "stats_parsed"])]
    #[case::json_and_struct(true, true, false, &["path", "stats", "stats_parsed"])]
    #[case::with_partition_values(
        true,
        true,
        true,
        &["path", "partitionValues", "partitionValues_parsed", "stats", "stats_parsed"])]
    #[case::partition_values_struct_disabled(
        true,
        false,
        true,
        &["path", "partitionValues", "stats"])]
    fn test_build_add_output_schema(
        #[case] write_stats_as_json: bool,
        #[case] write_stats_as_struct: bool,
        #[case] with_partition_schema: bool,
        #[case] expected_names: &[&str],
    ) {
        let config = StatsTransformConfig {
            write_stats_as_json,
            write_stats_as_struct,
        };
        let num_records = StructField::nullable(NUM_RECORDS, DataType::LONG);
        let stats_schema = Arc::new(StructType::new_unchecked([num_records]));
        let pv_schema = with_partition_schema.then(|| {
            Arc::new(StructType::new_unchecked([
                StructField::nullable("year", DataType::INTEGER),
                StructField::nullable("month", DataType::INTEGER),
            ]))
        });

        let (output_schema, _) =
            build_checkpoint_transform(&config, &stats_schema, pv_schema.as_ref());
        let add_field = output_schema
            .field(ADD_NAME)
            .expect("Expected output schema to contain add field");
        let DataType::Struct(add_schema) = add_field.data_type() else {
            panic!("Expected output add field to be a struct");
        };
        let field_names: Vec<_> = add_schema
            .fields()
            .map(|field| field.name.as_str())
            .collect();
        assert_eq!(field_names, expected_names);
    }
}