hamelin_analysis 0.23.5

//! Shared backward propagation for bounded and open-ended output requirements.

use chrono::Months;
use hamelin_lib::tree::{
    ast::clause::SortOrder,
    typed_ast::{
        command::TypedCommand,
        expression::{FieldReferenceResolution, FoldExpressionAlgebra, TypedFieldReference},
    },
};

use super::*;

/// Internal range with a genuinely open upper bound; no sentinel timestamp
/// participates in reverse evaluation or interval arithmetic.
#[derive(Clone, Copy, Debug)]
pub struct InputRange {
    pub start: Timestamp,
    pub end: Option<Timestamp>,
}

impl InputRange {
    pub fn union(self, other: Self) -> Self {
        Self {
            start: self.start.min(other.start),
            end: match (self.end, other.end) {
                (Some(a), Some(b)) => Some(a.max(b)),
                _ => None,
            },
        }
    }
}

pub fn compute_input_range_backward(
    pipeline: &TypedPipeline,
    requirements: &mut HashMap<Identifier, InputRange>,
    mut required: InputRange,
    field: &SimpleIdentifier,
    prevalidated_incremental: bool,
    definitions: &[TypedPipelineDef],
) -> Result<InputRange, IncrementalAnalysisError> {
    let commands = &pipeline
        .kind
        .as_ref()
        .map_err(|e| IncrementalAnalysisError::TreeHadError(e.as_ref().clone()))?
        .commands;
    let mut inputs: Option<InputRange> = None;
    let mut lineage: Option<Identifier> = None;
    for command in commands.iter().rev() {
        if let Some(path) = lineage.take() {
            let input = trace_field_path(command, &path)?;
            if input != field.clone().into() {
                lineage = Some(input);
            }
            continue;
        }
        match &command.kind {
            TypedCommandKind::Set(c) => {
                lineage = projected_field_path(&c.projections, field);
                if lineage.is_none() {
                    required = projection_input_range(&c.projections, required, field)?;
                }
            }
            TypedCommandKind::Select(c) => {
                if c.projections.lookup(&field.clone().into()).is_none() {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                }
                lineage = projected_field_path(&c.projections, field);
                if lineage.is_none() {
                    required = projection_input_range(&c.projections, required, field)?;
                }
            }
            TypedCommandKind::Window(c) => {
                // Interval expansion is supported only on the ascending timestamp
                // clock. Other expressions or directions need separate analysis.
                if c.sort_by.len() != 1
                    || c.sort_by[0].order != SortOrder::Asc
                    || !matches!(
                        &c.sort_by[0].expression.ast.kind,
                        ExpressionKind::FieldReference(reference)
                            if reference.field_name.valid_ref().ok() == Some(field)
                    )
                    || !c.sort_by[0].expression.fold(&mut OnlyTimestamp(field))
                {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                }
                required = projection_input_range(&c.projections, required, field)?;
                inputs = Some(inputs.map_or(required, |range| range.union(required)));
                let within = c
                    .within
                    .as_ref()
                    .ok_or(IncrementalAnalysisError::UnboundRange)?;
                check_expression_for_non_deterministic(within)?;
                let value = eval(within, &Environment::new())
                    .map_err(IncrementalAnalysisError::ExpressionEvaluationError)?;
                required = expand_window(required, value)?;
                inputs = Some(inputs.map_or(required, |range| range.union(required)));
                // Grouping assignments are evaluated before the window, so reverse
                // them after expanding the frame in the window's timestamp domain.
                required = projection_input_range(&c.group_by, required, field)?;
            }
            TypedCommandKind::Agg(c) => {
                if c.group_by.lookup(&field.clone().into()).is_none() {
                    return Err(IncrementalAnalysisError::AggWithoutTimestampGroupBy);
                }
                required = projection_input_range(&c.group_by, required, field)?;
            }
            TypedCommandKind::Distinct(c) => {
                if c.keys.lookup(&field.clone().into()).is_none() {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                }
                required = projection_input_range(&c.keys, required, field)?;
            }
            TypedCommandKind::Suppress(c) => {
                required = projection_input_range(&c.group_by, required, field)?;
                let ExpressionKind::IntervalLiteral(interval) = &c.interval.ast.kind else {
                    return Err(IncrementalAnalysisError::ValueNotInterval);
                };
                let (unit, multiplier) = interval
                    .trunc_unit_and_multiplier()
                    .ok_or(IncrementalAnalysisError::ValueNotInterval)?;
                required.start =
                    *truncate_timestamp(&TimestampValue::utc(required.start), &unit, multiplier)
                        .map_err(IncrementalAnalysisError::ExpressionEvaluationError)?
                        .instant();
                required.end = required
                    .end
                    .map(|end| {
                        let truncated =
                            truncate_timestamp(&TimestampValue::utc(end), &unit, multiplier)?;
                        next_truncation_boundary(&truncated, &unit, multiplier)
                            .map(|v| *v.instant())
                    })
                    .transpose()
                    .map_err(IncrementalAnalysisError::ExpressionEvaluationError)?;
            }
            TypedCommandKind::Drop(c) => {
                if c.dropped_fields.contains(&field.clone().into()) {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                }
            }
            TypedCommandKind::Explode(c) => {
                for item in &c.items {
                    check_expression_for_non_deterministic(&item.assignment.expression)?;
                    let target =
                        item.assignment.identifier.valid_ref().map_err(|e| {
                            IncrementalAnalysisError::TreeHadError(e.as_ref().clone())
                        })?;
                    if target.first() == *field {
                        return Err(IncrementalAnalysisError::TimestampLineageError);
                    }
                }
            }
            TypedCommandKind::Parse(c) => {
                if let Some(source) = &c.source {
                    check_expression_for_non_deterministic(source)?;
                }
                for identifier in &c.identifiers {
                    let target = identifier
                        .valid_ref()
                        .map_err(|e| IncrementalAnalysisError::TreeHadError(e.as_ref().clone()))?;
                    if target.first() == *field {
                        return Err(IncrementalAnalysisError::TimestampLineageError);
                    }
                }
            }
            TypedCommandKind::Unnest(c) => {
                check_expression_for_non_deterministic(&c.expression)?;
                let mut extracted = c.expression.resolved_type.as_ref();
                if let Type::Array(array) = extracted {
                    extracted = array.element_type.as_ref();
                }
                let Type::Struct(extracted) = extracted else {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                };
                if matches!(c.expression.resolved_type.as_ref(), Type::Struct(_))
                    && extracted.lookup(field).is_some()
                {
                    if let Some(source) = c.expression.unnest_source_drop_identifier() {
                        lineage = Some(source + field.clone().into());
                        continue;
                    }
                }
                if extracted.lookup(field).is_some()
                    || c.expression
                        .unnest_source_drop_identifier()
                        .is_some_and(|source| source.first() == *field)
                {
                    return Err(IncrementalAnalysisError::TimestampLineageError);
                }
            }
            TypedCommandKind::Where(_)
            | TypedCommandKind::Sort(_)
            | TypedCommandKind::Trimstrings(_)
            | TypedCommandKind::Within(_) => {}
            TypedCommandKind::From(TypedFromCommand { clauses })
            | TypedCommandKind::Union(TypedUnionCommand { clauses }) => {
                for reference in collect_static_refs(clauses)? {
                    if let Some(name) = dataset_ref_as_def_name(&reference) {
                        requirements
                            .entry(name)
                            .and_modify(|start| *start = start.union(required))
                            .or_insert(required);
                    }
                }
                return Ok(inputs.map_or(required, |range| range.union(required)));
            }
            TypedCommandKind::Join(TypedJoinCommand { right, condition })
            | TypedCommandKind::Lookup(TypedLookupCommand { right, condition }) => {
                if !prevalidated_incremental {
                    check_expression_for_non_deterministic(condition)?;
                    if right.alias.valid_ref().ok() == Some(field)
                        || !timestamp_free_table(right, field, definitions, &mut HashSet::new())?
                    {
                        return Err(IncrementalAnalysisError::CommandNotSupported(
                            command.ast.kind.command_name().to_string(),
                        ));
                    }
                }
            }
            TypedCommandKind::Match(c) if prevalidated_incremental => {
                if let Some(within) = &c.within {
                    let value = eval(within, &Environment::new())
                        .map_err(IncrementalAnalysisError::ExpressionEvaluationError)?;
                    let offset = match value {
                        Value::Interval(delta) => Value::Interval(-delta),
                        Value::CalendarInterval(months) => Value::CalendarInterval(
                            months
                                .checked_neg()
                                .ok_or(IncrementalAnalysisError::UnboundRange)?,
                        ),
                        _ => return Err(IncrementalAnalysisError::ValueNotInterval),
                    };
                    required.start = shift_timestamp(required.start, offset)?;
                }
                return Ok(inputs.map_or(required, |range| range.union(required)));
            }
            // Incremental callers have already checked command eligibility in
            // their forward pass (including lookup policy). The suffix API has
            // no such validation and must reject unsupported commands here.
            TypedCommandKind::Limit(_)
            | TypedCommandKind::Match(_)
            | TypedCommandKind::Nest(_)
            | TypedCommandKind::Append(_) => {
                if !prevalidated_incremental {
                    return Err(IncrementalAnalysisError::CommandNotSupported(
                        command.ast.kind.command_name().to_string(),
                    ));
                }
            }
            TypedCommandKind::Error(error) => {
                return Err(IncrementalAnalysisError::TreeHadError(
                    error.as_ref().clone(),
                ));
            }
        }
        inputs = Some(inputs.map_or(required, |range| range.union(required)));
    }
    if lineage.is_some() {
        return Err(IncrementalAnalysisError::TimestampLineageError);
    }
    Err(IncrementalAnalysisError::EmptyPipeline)
}

fn projected_field_path(projections: &Projections, field: &SimpleIdentifier) -> Option<Identifier> {
    let projection = projections.lookup(&field.clone().into())?;
    let path = projection.expression.unnest_source_drop_identifier()?;
    (path != field.clone().into()).then_some(path)
}

fn trace_field_path(
    command: &TypedCommand,
    path: &Identifier,
) -> Result<Identifier, IncrementalAnalysisError> {
    match &command.kind {
        TypedCommandKind::Nest(c) => {
            let prefix = c
                .identifier
                .valid_ref()
                .map_err(|e| IncrementalAnalysisError::TreeHadError(e.as_ref().clone()))?;
            let parts = path.segments();
            let prefix = prefix.segments();
            if !parts.starts_with(prefix) || parts.len() <= prefix.len() {
                return Err(IncrementalAnalysisError::TimestampLineageError);
            }
            Ok(parts[prefix.len() + 1..].iter().fold(
                parts[prefix.len()].clone().into(),
                |path: Identifier, part| path + part.clone().into(),
            ))
        }
        TypedCommandKind::Set(c) => trace_projection_path(&c.projections, path, false),
        TypedCommandKind::Select(c) => trace_projection_path(&c.projections, path, true),
        TypedCommandKind::Unnest(c) => {
            let Type::Struct(fields) = c.expression.resolved_type.as_ref() else {
                return Err(IncrementalAnalysisError::TimestampLineageError);
            };
            if fields.lookup(&path.first()).is_some() {
                let source = c
                    .expression
                    .unnest_source_drop_identifier()
                    .ok_or(IncrementalAnalysisError::TimestampLineageError)?;
                Ok(source + path.clone())
            } else if c
                .expression
                .unnest_source_drop_identifier()
                .is_some_and(|source| path.segments().starts_with(source.segments()))
            {
                Err(IncrementalAnalysisError::TimestampLineageError)
            } else {
                Ok(path.clone())
            }
        }
        TypedCommandKind::Drop(c) => {
            if c.dropped_fields
                .iter()
                .any(|dropped| path.segments().starts_with(dropped.segments()))
            {
                return Err(IncrementalAnalysisError::TimestampLineageError);
            }
            Ok(path.clone())
        }
        TypedCommandKind::Where(_)
        | TypedCommandKind::Sort(_)
        | TypedCommandKind::Within(_)
        | TypedCommandKind::Trimstrings(_) => Ok(path.clone()),
        _ => Err(IncrementalAnalysisError::TimestampLineageError),
    }
}

fn trace_projection_path(
    projections: &Projections,
    path: &Identifier,
    replaces_schema: bool,
) -> Result<Identifier, IncrementalAnalysisError> {
    for assignment in &projections.assignments {
        let target = assignment
            .identifier
            .valid_ref()
            .map_err(|e| IncrementalAnalysisError::TreeHadError(e.as_ref().clone()))?;
        if path.segments().starts_with(target.segments()) {
            let source = assignment
                .expression
                .unnest_source_drop_identifier()
                .ok_or(IncrementalAnalysisError::TimestampLineageError)?;
            return Ok(path.segments()[target.segments().len()..]
                .iter()
                .fold(source, |path, part| path + part.clone().into()));
        }
    }
    if replaces_schema {
        return Err(IncrementalAnalysisError::TimestampLineageError);
    }
    Ok(path.clone())
}

// DEF projections can hide a source timestamp, so the final schema is insufficient.
fn timestamp_free_table(
    table: &TypedTableAlias,
    field: &SimpleIdentifier,
    definitions: &[TypedPipelineDef],
    visited: &mut HashSet<Identifier>,
) -> Result<bool, IncrementalAnalysisError> {
    if let Some(error) = table.resolved.error() {
        return Err(IncrementalAnalysisError::TreeHadError(
            error.as_ref().clone(),
        ));
    }
    if table
        .resolved
        .environment()
        .lookup(&field.clone().into())
        .is_some()
    {
        return Ok(false);
    }
    timestamp_free_definition(&rhs_dataset_ref(table)?, field, definitions, visited)
}

fn timestamp_free_definition(
    reference: &DatasetIdentifier,
    field: &SimpleIdentifier,
    definitions: &[TypedPipelineDef],
    visited: &mut HashSet<Identifier>,
) -> Result<bool, IncrementalAnalysisError> {
    let Some(name) = dataset_ref_as_def_name(reference) else {
        return Ok(true);
    };
    let Some(definition) = find_pipeline_def_by_name(&name, definitions) else {
        return Ok(true);
    };
    if !visited.insert(name) {
        return Ok(true);
    }
    let commands = &definition
        .pipeline
        .kind
        .as_ref()
        .map_err(|error| IncrementalAnalysisError::TreeHadError(error.as_ref().clone()))?
        .commands;
    for command in commands {
        if command.input_schema.lookup(&field.clone().into()).is_some()
            || command
                .output_schema
                .lookup(&field.clone().into())
                .is_some()
        {
            return Ok(false);
        }
        match &command.kind {
            TypedCommandKind::From(TypedFromCommand { clauses })
            | TypedCommandKind::Union(TypedUnionCommand { clauses }) => {
                for clause in clauses {
                    if clause.environment().lookup(&field.clone().into()).is_some() {
                        return Ok(false);
                    }
                }
                for reference in collect_static_refs(clauses)? {
                    if !timestamp_free_definition(&reference, field, definitions, visited)? {
                        return Ok(false);
                    }
                }
            }
            TypedCommandKind::Join(TypedJoinCommand { right, .. })
            | TypedCommandKind::Lookup(TypedLookupCommand { right, .. }) => {
                if !timestamp_free_table(right, field, definitions, visited)? {
                    return Ok(false);
                }
            }
            TypedCommandKind::Match(_) => return Ok(false),
            TypedCommandKind::Error(error) => {
                return Err(IncrementalAnalysisError::TreeHadError(
                    error.as_ref().clone(),
                ));
            }
            _ => {}
        }
    }
    Ok(true)
}

fn projection_input_range(
    projections: &Projections,
    output: InputRange,
    field: &SimpleIdentifier,
) -> Result<InputRange, IncrementalAnalysisError> {
    let Some(projection) = projections.lookup(&field.clone().into()) else {
        return Ok(output);
    };
    check_expression_for_non_deterministic(&projection.expression)?;
    if !projection.expression.fold(&mut OnlyTimestamp(field)) {
        return Err(IncrementalAnalysisError::TimestampLineageError);
    }
    let constraint = Constraint::Range {
        min: Some(TimestampValue::utc(output.start).into()),
        max: output.end.map(|end| TimestampValue::utc(end).into()),
    };
    match reverse_eval(&projection.expression, constraint, &Environment::new()) {
        Ok(Some(Constraint::Range {
            min: Some(Value::Timestamp(start)),
            max,
        })) => {
            let end = match max {
                Some(Value::Timestamp(end)) if output.end.is_some() => Some(*end.instant()),
                None if output.end.is_none() => None,
                _ => return Err(IncrementalAnalysisError::TimestampLineageError),
            };
            Ok(InputRange {
                start: *start.instant(),
                end,
            })
        }
        _ => Err(IncrementalAnalysisError::TimestampLineageError),
    }
}

fn expand_window(range: InputRange, value: Value) -> Result<InputRange, IncrementalAnalysisError> {
    let (lower, upper) = match value {
        Value::Range(value) => (value.lower.clone(), value.upper.clone()),
        value => (Some(value.clone()), Some(value)),
    };
    let start = range.start.min(shift_timestamp(
        range.start,
        lower.ok_or(IncrementalAnalysisError::UnboundRange)?,
    )?);
    let end = range
        .end
        .map(|end| {
            shift_timestamp(end, upper.ok_or(IncrementalAnalysisError::UnboundRange)?)
                .map(|shifted| end.max(shifted))
        })
        .transpose()?;
    Ok(InputRange { start, end })
}

struct OnlyTimestamp<'a>(&'a SimpleIdentifier);
impl FoldExpressionAlgebra<bool> for OnlyTimestamp<'_> {
    fn combine(&mut self, children: impl IntoIterator<Item = bool>) -> bool {
        children.into_iter().all(|v| v)
    }
    fn field_reference(&mut self, node: &TypedFieldReference, _: &TypedExpression) -> bool {
        matches!(node.resolution, FieldReferenceResolution::Pipeline)
            && node.field_name.valid_ref().ok() == Some(self.0)
    }
    fn leaf(&mut self, expr: &TypedExpression) -> bool {
        match &expr.ast.kind {
            ExpressionKind::FieldReference(reference) => {
                reference.field_name.valid_ref().ok() == Some(self.0)
            }
            _ => true,
        }
    }
}

fn shift_timestamp(start: Timestamp, value: Value) -> Result<Timestamp, IncrementalAnalysisError> {
    match value {
        Value::Interval(delta) => start
            .checked_add_signed(delta)
            .ok_or(IncrementalAnalysisError::UnboundRange),
        Value::CalendarInterval(months) => {
            let offset = Months::new(months.unsigned_abs());
            if months < 0 {
                start.checked_sub_months(offset)
            } else {
                start.checked_add_months(offset)
            }
            .ok_or(IncrementalAnalysisError::UnboundRange)
        }
        _ => Err(IncrementalAnalysisError::ValueNotInterval),
    }
}