delta_kernel 0.30.0

//! Expression handling based on arrow-rs compute kernels.
use std::sync::Arc;

pub(crate) use evaluate_expression::extract_column;
use evaluate_expression::{evaluate_expression, evaluate_predicate};
use tracing::debug;

use super::arrow_conversion::{TryFromKernel as _, TryIntoArrow as _};
use crate::arrow::array::{self, ArrayBuilder, ArrayRef, RecordBatch, StructArray};
use crate::arrow::datatypes::{
    DataType as ArrowDataType, Field as ArrowField, Schema as ArrowSchema, TimeUnit,
};
use crate::engine::arrow_data::{extract_record_batch, ArrowEngineData};
use crate::engine::arrow_utils::apply_schema::{apply_schema, apply_schema_to};
use crate::error::{KernelError, KernelResult};
use crate::expressions::{ArrayData, Expression, ExpressionRef, PredicateRef, Scalar};
use crate::schema::{DataType, PrimitiveType, SchemaRef};
use crate::utils::require;
use crate::{EngineData, EvaluationHandler, ExpressionEvaluator, PredicateEvaluator};

pub mod evaluate_expression;
pub mod opaque;
mod timestamp_timezone;

#[cfg(test)]
mod tests;

// TODO leverage scalars / Datum

impl Scalar {
    /// Convert scalar to arrow array.
    pub fn to_array(&self, num_rows: usize) -> KernelResult<ArrayRef> {
        let data_type = ArrowDataType::try_from_kernel(&self.data_type())?;
        let mut builder = array::make_builder(&data_type, num_rows);
        self.append_to(&mut builder, num_rows)?;
        Ok(builder.finish())
    }

    // Arrow uses composable "builders" to assemble arrays one row at a time. Each concrete `Array`
    // type has a corresponding concrete `ArrayBuilder` type. For primitive types, the builder just
    // needs to `append` one value per row. For complex types, the builder needs to recursively
    // append values to each of its children as needed, and then its own `append` only defines the
    // validity for the row. Unfortunately, there is no generic way to append values to builders;
    // the `ArrayBuilder` trait only knows how to `finalize` itself to produce an `ArrayRef`. So we
    // have to cast each builder to the appropriate type, based on the scalar's data type. For
    // details, refer to the arrow documentation:
    //
    // https://docs.rs/arrow/latest/arrow/array/struct.PrimitiveBuilder.html
    // https://docs.rs/arrow/latest/arrow/array/struct.GenericListBuilder.html
    // https://docs.rs/arrow/latest/arrow/array/struct.StructBuilder.html
    //
    // NOTE: `ListBuilder` and `MapBuilder` are take generic element/key/value builders in order to
    // work with specific builder types directly. However, `array::make_builder` instantiates them
    // with `Box<dyn Builder>` instead, which greatly simplifies our job in working with them. We
    // can just extract the builder trait, and let recursive calls cast it to the desired type.
    //
    // WARNING: List and map builders do _NOT_ require appending any child entries to NULL list/map
    // rows, because empty list/map is a valid state. But struct builders _DO_ require appending
    // (possibly NULL) entries in order to preserve consistent row counts between the struct and its
    // fields.
    fn append_to(&self, builder: &mut dyn ArrayBuilder, num_rows: usize) -> KernelResult<()> {
        use Scalar::*;
        macro_rules! builder_as {
            ($t:ty) => {{
                builder.as_any_mut().downcast_mut::<$t>().ok_or_else(|| {
                    KernelError::invalid_expression(format!(
                        "Invalid builder for {}",
                        self.data_type()
                    ))
                })?
            }};
        }

        // Use append_value_n for primitive builders that support batch append
        macro_rules! append_val_n_as {
            ($t:ty, $val:expr) => {{
                let builder = builder_as!($t);
                builder.append_value_n($val, num_rows);
            }};
        }

        // Use append_value in a loop for builders without batch append (String, Binary)
        // TODO: Remove after https://github.com/apache/arrow-rs/pull/9426 gets in
        macro_rules! append_val_as {
            ($t:ty, $val:expr) => {{
                let builder = builder_as!($t);
                for _ in 0..num_rows {
                    builder.append_value($val);
                }
            }};
        }

        match self {
            Integer(val) => append_val_n_as!(array::Int32Builder, *val),
            Long(val) => append_val_n_as!(array::Int64Builder, *val),
            Short(val) => append_val_n_as!(array::Int16Builder, *val),
            Byte(val) => append_val_n_as!(array::Int8Builder, *val),
            Float(val) => append_val_n_as!(array::Float32Builder, *val),
            Double(val) => append_val_n_as!(array::Float64Builder, *val),
            String(val) => append_val_as!(array::StringBuilder, val),
            Boolean(val) => builder_as!(array::BooleanBuilder).append_n(num_rows, *val),
            Timestamp(val) | TimestampNtz(val) => {
                // timezone was already set at builder construction time
                append_val_n_as!(array::TimestampMicrosecondBuilder, *val)
            }
            IntervalYearMonth(val) => append_val_n_as!(array::Int32Builder, *val),
            IntervalDayTime(val) => append_val_n_as!(array::Int64Builder, *val),
            Date(val) => append_val_n_as!(array::Date32Builder, *val),
            Binary(val) => append_val_as!(array::BinaryBuilder, val),
            // precision and scale were already set at builder construction time
            Decimal(val) => append_val_n_as!(array::Decimal128Builder, val.bits()),
            Struct(data) => {
                let builder = builder_as!(array::StructBuilder);
                require!(
                    builder.num_fields() == data.fields().len(),
                    KernelError::generic("Struct builder has wrong number of fields")
                );
                let field_builders = builder.field_builders_mut().iter_mut();
                for (builder, value) in field_builders.zip(data.values()) {
                    value.append_to(builder, num_rows)?;
                }
                // TODO: Loop can be removed after: https://github.com/apache/arrow-rs/pull/9430
                for _ in 0..num_rows {
                    builder.append(true);
                }
            }
            Array(data) => {
                let builder = builder_as!(array::ListBuilder<Box<dyn ArrayBuilder>>);
                for _ in 0..num_rows {
                    for value in data.array_elements() {
                        value.append_to(builder.values(), 1)?;
                    }
                    builder.append(true);
                }
            }
            Map(data) => {
                let builder =
                    builder_as!(array::MapBuilder<Box<dyn ArrayBuilder>, Box<dyn ArrayBuilder>>);
                for _ in 0..num_rows {
                    for (key, val) in data.pairs() {
                        key.append_to(builder.keys(), 1)?;
                        val.append_to(builder.values(), 1)?;
                    }
                    builder.append(true)?;
                }
            }
            Null(data_type) => Self::append_null(builder, data_type, num_rows)?,
        }

        Ok(())
    }

    fn append_null(
        builder: &mut dyn ArrayBuilder,
        data_type: &DataType,
        num_rows: usize,
    ) -> KernelResult<()> {
        // Almost the same as above -- differs only in the data type parameter
        macro_rules! builder_as {
            ($t:ty) => {{
                builder.as_any_mut().downcast_mut::<$t>().ok_or_else(|| {
                    KernelError::invalid_expression(format!("Invalid builder for {data_type}"))
                })?
            }};
        }

        macro_rules! append_nulls_as {
            ($t:ty) => {{
                let builder = builder_as!($t);
                builder.append_nulls(num_rows);
            }};
        }

        match *data_type {
            #[cfg(feature = "udt-in-dev")]
            DataType::UserDefined(_) => {
                Self::append_null(builder, data_type.physical_type(), num_rows)?
            }
            DataType::INTEGER => append_nulls_as!(array::Int32Builder),
            DataType::LONG => append_nulls_as!(array::Int64Builder),
            DataType::SHORT => append_nulls_as!(array::Int16Builder),
            DataType::BYTE => append_nulls_as!(array::Int8Builder),
            DataType::FLOAT => append_nulls_as!(array::Float32Builder),
            DataType::DOUBLE => append_nulls_as!(array::Float64Builder),
            DataType::STRING => append_nulls_as!(array::StringBuilder),
            DataType::BOOLEAN => append_nulls_as!(array::BooleanBuilder),
            DataType::TIMESTAMP | DataType::TIMESTAMP_NTZ => {
                append_nulls_as!(array::TimestampMicrosecondBuilder)
            }
            DataType::DATE => append_nulls_as!(array::Date32Builder),
            DataType::BINARY => append_nulls_as!(array::BinaryBuilder),
            DataType::Primitive(PrimitiveType::Decimal(_)) => {
                append_nulls_as!(array::Decimal128Builder)
            }
            // A variant is physically a struct (`metadata`/`value`, plus any shredded fields), so a
            // null variant is a null struct and builds through the same StructBuilder path. (Only
            // the null case is reachable: there is no non-null `Scalar::Variant`.)
            DataType::Struct(ref stype) | DataType::Variant(ref stype) => {
                // WARNING: Unlike ArrayBuilder and MapBuilder, StructBuilder always requires us to
                // insert an entry for each child builder, even when we're inserting NULL.
                let builder = builder_as!(array::StructBuilder);
                require!(
                    builder.num_fields() == stype.num_fields(),
                    KernelError::generic("Struct builder has wrong number of fields")
                );
                let field_builders = builder.field_builders_mut().iter_mut();
                for (builder, field) in field_builders.zip(stype.fields()) {
                    Self::append_null(builder, &field.data_type, num_rows)?;
                }
                builder.append_nulls(num_rows);
            }
            DataType::Array(_) => append_nulls_as!(array::ListBuilder<Box<dyn ArrayBuilder>>),
            DataType::Map(_) => {
                // For some reason, there is no `MapBuilder::append_null` method -- even tho
                // StructBuilder and ListBuilder both provide it.
                let builder =
                    builder_as!(array::MapBuilder<Box<dyn ArrayBuilder>, Box<dyn ArrayBuilder>>);
                // TODO: Can be removed after https://github.com/apache/arrow-rs/pull/9432
                for _ in 0..num_rows {
                    builder.append(false)?;
                }
            }
            DataType::VOID => append_nulls_as!(array::NullBuilder),
            // Intervals are exposed as their physical integer (i32 months / i64 microseconds).
            DataType::INTERVAL_YEAR_MONTH => append_nulls_as!(array::Int32Builder),
            DataType::INTERVAL_DAY_TIME => append_nulls_as!(array::Int64Builder),
            #[cfg(feature = "geo-type-in-dev")]
            DataType::Primitive(PrimitiveType::Geometry(_) | PrimitiveType::Geography(_)) => {
                return Err(KernelError::unsupported(
                    "Geo is not supported as scalar yet.",
                ));
            }
        }
        Ok(())
    }
}

impl ArrayData {
    /// Convert kernel [`ArrayData`] to an Arrow [`ArrayRef`] of the equivalent type.
    pub fn to_arrow(&self) -> KernelResult<ArrayRef> {
        let arrow_data_type = ArrowDataType::try_from_kernel(self.array_type().element_type())?;

        let elements = self.array_elements();
        let mut builder = array::make_builder(&arrow_data_type, elements.len());
        for element in elements {
            element.append_to(&mut builder, 1)?;
        }

        Ok(builder.finish())
    }
}

#[derive(Debug)]
pub struct ArrowEvaluationHandler;

impl EvaluationHandler for ArrowEvaluationHandler {
    fn new_expression_evaluator(
        &self,
        schema: SchemaRef,
        expression: ExpressionRef,
        output_type: DataType,
    ) -> KernelResult<Arc<dyn ExpressionEvaluator>> {
        Ok(Arc::new(DefaultExpressionEvaluator {
            input_schema: schema,
            expression,
            output_type,
        }))
    }

    fn new_predicate_evaluator(
        &self,
        schema: SchemaRef,
        predicate: PredicateRef,
    ) -> KernelResult<Arc<dyn PredicateEvaluator>> {
        Ok(Arc::new(DefaultPredicateEvaluator {
            input_schema: schema,
            predicate,
        }))
    }

    fn create_many(
        &self,
        schema: SchemaRef,
        rows: Vec<Vec<Scalar>>,
    ) -> KernelResult<Box<dyn EngineData>> {
        let arrow_schema: Arc<ArrowSchema> = Arc::new(schema.as_ref().try_into_arrow()?);
        if rows.is_empty() {
            return Ok(Box::new(ArrowEngineData::new(RecordBatch::new_empty(
                arrow_schema,
            ))));
        }

        let num_rows = rows.len();
        let num_fields = schema.fields().len();
        for (row_idx, row) in rows.iter().enumerate() {
            if row.len() != num_fields {
                return Err(KernelError::generic(format!(
                    "Row {} has {} scalars but schema has {} fields",
                    row_idx,
                    row.len(),
                    num_fields
                )));
            }
        }

        let mut builders: Vec<Box<dyn ArrayBuilder>> = arrow_schema
            .fields()
            .iter()
            .map(|field| array::make_builder(field.data_type(), num_rows))
            .collect();

        let fields: Vec<_> = schema.fields().collect();
        for (col_idx, builder) in builders.iter_mut().enumerate() {
            let field_name = fields[col_idx].name();
            for (row_idx, row) in rows.iter().enumerate() {
                row[col_idx].append_to(builder.as_mut(), 1).map_err(|e| {
                    KernelError::generic(format!(
                        "Row {row_idx}, field '{field_name}' \
                            (expected type {}, got {}): {e}",
                        fields[col_idx].data_type(),
                        row[col_idx].data_type()
                    ))
                })?;
            }
        }

        let arrays: Vec<ArrayRef> = builders.into_iter().map(|mut b| b.finish()).collect();

        Ok(Box::new(ArrowEngineData::new(RecordBatch::try_new(
            arrow_schema,
            arrays,
        )?)))
    }
}

#[derive(Debug)]
pub struct DefaultExpressionEvaluator {
    input_schema: SchemaRef,
    expression: ExpressionRef,
    output_type: DataType,
}

impl ExpressionEvaluator for DefaultExpressionEvaluator {
    fn evaluate(&self, batch: &dyn EngineData) -> KernelResult<Box<dyn EngineData>> {
        debug!("Arrow evaluator evaluating: {:#?}", self.expression);
        let batch = extract_record_batch(batch)?;
        // TODO(#3263): Validate nested fields.
        validate_data_schema_top_level(&self.input_schema, batch.schema().as_ref())?;
        let batch = match (self.expression.as_ref(), &self.output_type) {
            (Expression::StructPatch(patch), DataType::Struct(_)) if patch.is_empty() => {
                // Empty patch optimization: Skip expression evaluation and directly apply the
                // output schema to the input RecordBatch. This is used to cheaply apply a new
                // output schema to existing data without changing it, e.g. for column mapping.
                let array = match patch.input_path() {
                    None => Arc::new(StructArray::from(batch.clone())),
                    Some(path) => extract_column(batch, path)?,
                };
                apply_schema(&array, &self.output_type)?
            }
            (expr, output_type @ DataType::Struct(_)) => {
                let array_ref = evaluate_expression(expr, batch, Some(output_type))?;
                apply_schema(&array_ref, output_type)?
            }
            (expr, output_type) => {
                let array_ref = evaluate_expression(expr, batch, Some(output_type))?;
                let array_ref = apply_schema_to(&array_ref, output_type)?;
                let arrow_type = ArrowDataType::try_from_kernel(output_type)?;
                let schema = ArrowSchema::new(vec![ArrowField::new("output", arrow_type, true)]);
                RecordBatch::try_new(Arc::new(schema), vec![array_ref])?
            }
        };

        Ok(Box::new(ArrowEngineData::new(batch)))
    }
}

#[derive(Debug)]
pub struct DefaultPredicateEvaluator {
    input_schema: SchemaRef,
    predicate: PredicateRef,
}

impl PredicateEvaluator for DefaultPredicateEvaluator {
    fn evaluate(&self, batch: &dyn EngineData) -> KernelResult<Box<dyn EngineData>> {
        debug!("Arrow evaluator evaluating: {:#?}", self.predicate);
        let batch = extract_record_batch(batch)?;
        // TODO(#3263): Validate nested fields.
        validate_data_schema_top_level(&self.input_schema, batch.schema().as_ref())?;
        let array = evaluate_predicate(&self.predicate, batch, false)?;
        let schema = ArrowSchema::new(vec![ArrowField::new(
            "output",
            ArrowDataType::Boolean,
            true,
        )]);
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)])?;
        Ok(Box::new(ArrowEngineData::new(batch)))
    }
}
/// Validates that each expected field exists and has a compatible type at top-level.
fn validate_data_schema_top_level(
    expected_schema: &SchemaRef,
    data_schema: &ArrowSchema,
) -> KernelResult<()> {
    let mut data_fields = data_schema.fields().iter();
    // Some Kernel code does not provide the full input schema to the evaluator. For example,
    // `scan_metadata_from` may evaluate scan rows containing optional `stats_parsed` and
    // `partitionValues_parsed` columns using only the base scan-row schema. As a result,
    // we allow `data_schema` to contain extra fields.
    // TODO(#3263): Require evaluator input schemas to declare every top-level field.
    for expected_field in expected_schema.fields() {
        let data_field = data_fields
            .find(|field| field.name() == expected_field.name())
            .ok_or_else(|| {
                let mismatch = if data_schema
                    .fields()
                    .iter()
                    .any(|field| field.name() == expected_field.name())
                {
                    "out of order"
                } else {
                    "missing"
                };
                KernelError::schema(format!(
                    "Expected schema field '{}' is {mismatch} in data schema fields {:?}",
                    expected_field.name(),
                    data_schema
                        .fields()
                        .iter()
                        .map(|field| field.name())
                        .collect::<Vec<_>>()
                ))
            })?;
        require!(
            top_level_types_compatible(expected_field.data_type(), data_field.data_type()),
            KernelError::schema(format!(
                "Expected schema type for '{}' does not match the data schema type: {:?} != {:?}",
                expected_field.name(),
                expected_field.data_type(),
                data_field.data_type()
            ))
        );
    }
    Ok(())
}

/// Checks top-level type compatibility using the Arrow-to-Kernel mappings from
/// [`TryIntoKernel`](super::arrow_conversion::TryIntoKernel).
///
/// Unlike a full conversion, this does not inspect nested types or field metadata.
fn top_level_types_compatible(expected_type: &DataType, data_type: &ArrowDataType) -> bool {
    let expected_type = expected_type.physical_type();
    match (expected_type, data_type) {
        // Dictionary types have the same logical type as their values.
        (_, ArrowDataType::Dictionary(_, value_type)) => {
            top_level_types_compatible(expected_type, value_type)
        }
        (DataType::Primitive(expected), data_type) => {
            primitive_types_compatible(expected, data_type)
        }
        (DataType::Struct(_), ArrowDataType::Struct(_)) => true,
        (
            DataType::Array(_),
            ArrowDataType::List(_)
            | ArrowDataType::ListView(_)
            | ArrowDataType::LargeList(_)
            | ArrowDataType::LargeListView(_)
            | ArrowDataType::FixedSizeList(_, _),
        ) => true,
        (DataType::Map(_), ArrowDataType::Map(_, _)) => true,
        // Arrow has no Variant type, and it will be converted to structs.
        (DataType::Variant(_), ArrowDataType::Struct(_)) => true,
        _ => false,
    }
}

fn primitive_types_compatible(expected: &PrimitiveType, data_type: &ArrowDataType) -> bool {
    match (expected, data_type) {
        (
            PrimitiveType::String,
            ArrowDataType::Utf8 | ArrowDataType::LargeUtf8 | ArrowDataType::Utf8View,
        ) => true,
        (PrimitiveType::Long, ArrowDataType::Int64 | ArrowDataType::UInt64) => true,
        (PrimitiveType::Integer, ArrowDataType::Int32 | ArrowDataType::UInt32) => true,
        (PrimitiveType::Short, ArrowDataType::Int16 | ArrowDataType::UInt16) => true,
        (PrimitiveType::Byte, ArrowDataType::Int8 | ArrowDataType::UInt8) => true,
        (PrimitiveType::Float, ArrowDataType::Float32) => true,
        (PrimitiveType::Double, ArrowDataType::Float64) => true,
        (PrimitiveType::Boolean, ArrowDataType::Boolean) => true,
        (
            PrimitiveType::Binary,
            ArrowDataType::Binary
            | ArrowDataType::FixedSizeBinary(_)
            | ArrowDataType::LargeBinary
            | ArrowDataType::BinaryView,
        ) => true,
        (PrimitiveType::Decimal(expected), ArrowDataType::Decimal128(precision, scale)) => {
            *precision == expected.precision()
                && u8::try_from(*scale).is_ok_and(|scale| scale == expected.scale())
        }
        (PrimitiveType::Date, ArrowDataType::Date32 | ArrowDataType::Date64) => true,
        (
            PrimitiveType::Timestamp,
            ArrowDataType::Timestamp(
                TimeUnit::Millisecond | TimeUnit::Microsecond | TimeUnit::Nanosecond,
                Some(timezone),
            ),
        ) => timezone.eq_ignore_ascii_case("utc"),
        (
            PrimitiveType::TimestampNtz,
            ArrowDataType::Timestamp(
                TimeUnit::Millisecond | TimeUnit::Microsecond | TimeUnit::Nanosecond,
                None,
            ),
        ) => true,
        (PrimitiveType::Void, ArrowDataType::Null) => true,
        (PrimitiveType::IntervalYearMonth, ArrowDataType::Int32 | ArrowDataType::UInt32) => true,
        (PrimitiveType::IntervalDayTime, ArrowDataType::Int64 | ArrowDataType::UInt64) => true,
        _ => false,
    }
}
#[cfg(test)]
fn expected_timestamp_micros(value: &str) -> i64 {
    chrono::DateTime::parse_from_rfc3339(value)
        .unwrap()
        .timestamp_micros()
}