algeon-datafusion 0.1.2

//! Name bindings and parent demand shared by every lowering step.
//!
//! DataFusion addresses columns by index and allows duplicate names; the
//! native IR addresses them by unique name. `FieldBinding` records the mapping
//! per lowered node so expressions above resolve DataFusion indices to native
//! names, and `frontend_schema` restores DataFusion names at the host sink.

use std::{
    collections::{BTreeMap, BTreeSet},
    sync::Arc,
};

use arrow::datatypes::{Field, Schema, SchemaRef};
use datafusion::physical_expr::expressions::Column;
use query_planning::plan::NodeId as NativeNodeId;

use super::naming::{sanitize_field_name, unique_name_with_suffix};
use crate::capability::CapabilityReason;

/// Native node plus the schema visible to its parent lowering step.
#[derive(Debug)]
pub(super) struct LoweredNativeInput {
    pub(super) node: NativeNodeId,
    /// Executable schema carried by the native query engine. Field names in
    /// this schema are unique within the node output.
    pub(super) schema: SchemaRef,
    /// Adapter-side mapping from the DataFusion physical output schema to the
    /// executable native schema above.
    pub(super) bindings: Vec<FieldBinding>,
    /// The node emits row-count weights in place of the input's rows, which
    /// only a parent that demanded the row count alone reads.
    pub(super) row_count_weights: bool,
}

impl LoweredNativeInput {
    pub(super) fn with_bindings(node: NativeNodeId, schema: SchemaRef, bindings: Vec<FieldBinding>) -> Self {
        Self { node, schema, bindings, row_count_weights: false }
    }

    pub(super) fn row_count_weights(node: NativeNodeId, schema: SchemaRef) -> Self {
        Self { node, schema, bindings: Vec::new(), row_count_weights: true }
    }

    pub(super) fn binding_for_column(&self, column: &Column) -> Result<&FieldBinding, CapabilityReason> {
        self.binding_for_frontend_index(column.index())
    }

    pub(super) fn binding_for_frontend_index(&self, frontend_index: usize) -> Result<&FieldBinding, CapabilityReason> {
        self.bindings
            .iter()
            .find(|binding| binding.frontend_index == frontend_index)
            .ok_or(CapabilityReason::UnsupportedExpression)
    }

    pub(super) fn frontend_schema(&self) -> Result<SchemaRef, CapabilityReason> {
        if self.schema.fields().len() != self.bindings.len() {
            return Err(CapabilityReason::UnsupportedExpression);
        }
        let fields = self
            .schema
            .fields()
            .iter()
            .zip(self.bindings.iter())
            .map(|(field, binding)| rename_field(field.as_ref(), binding.frontend_name.as_str()))
            .collect::<Vec<_>>();
        Ok(Arc::new(Schema::new(fields)))
    }
}

/// Mapping from a DataFusion output field to its executable native field name.
#[derive(Debug, Clone)]
pub(super) struct FieldBinding {
    pub(super) frontend_index: usize,
    pub(super) frontend_name: String,
    pub(super) native_name: String,
}

/// What a parent reads from the input it lowers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(in crate::native) struct Demand {
    /// Output fields the parent reads.
    pub(super) fields: RequiredFields,
    /// The parent reads only how many rows the input has. A Parquet scan that
    /// needs no column to produce its rows answers with row-count weights.
    pub(super) row_count_only: bool,
    /// The parent keeps at most this many rows of the input, and any rows
    /// satisfy it.
    pub(super) row_limit: Option<usize>,
    /// The parent selects `fields` from its input itself, so the input may
    /// also emit fields the parent does not read. This avoids adding a
    /// projection that only drops them.
    pub(super) extra_fields_allowed: bool,
    /// The parent reads the input's rows in the order DataFusion reports for
    /// the input. A scan of several sorted runs, which native lowering reads
    /// one after another, does not reproduce that order.
    pub(in crate::native) ordered: bool,
}

impl Demand {
    pub(super) fn all() -> Self {
        Self::fields(RequiredFields::all())
    }

    pub(super) fn fields(fields: RequiredFields) -> Self {
        Self { fields, row_count_only: false, row_limit: None, extra_fields_allowed: false, ordered: false }
    }
}

/// Frontend field set required by a parent plan fragment.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) enum RequiredFields {
    All,
    Only { indices: BTreeSet<usize>, predicate_indices: BTreeSet<usize> },
}

impl RequiredFields {
    pub(super) fn all() -> Self {
        Self::All
    }

    pub(super) fn only(required: BTreeSet<usize>) -> Self {
        Self::Only { indices: required, predicate_indices: BTreeSet::new() }
    }

    pub(super) fn only_with_predicates(mut required: BTreeSet<usize>, predicate_indices: BTreeSet<usize>) -> Self {
        required.extend(predicate_indices.iter().copied());
        Self::Only { indices: required, predicate_indices }
    }

    pub(super) fn is_all(&self) -> bool {
        matches!(self, Self::All)
    }

    pub(super) fn contains(&self, index: usize) -> bool {
        match self {
            Self::All => true,
            Self::Only { indices, .. } => indices.contains(&index),
        }
    }

    pub(super) fn requested_indices(&self, field_count: usize) -> Result<Vec<usize>, CapabilityReason> {
        match self {
            Self::All => Ok((0..field_count).collect()),
            Self::Only { indices, .. } => {
                if indices.iter().any(|index| *index >= field_count) {
                    return Err(CapabilityReason::UnsupportedExpression);
                }
                Ok(indices.iter().copied().collect())
            },
        }
    }

    pub(super) fn predicate_indices(&self, field_count: usize) -> Result<Vec<usize>, CapabilityReason> {
        match self {
            Self::All => Ok(Vec::new()),
            Self::Only { predicate_indices, .. } => {
                if predicate_indices.iter().any(|index| *index >= field_count) {
                    return Err(CapabilityReason::UnsupportedExpression);
                }
                Ok(predicate_indices.iter().copied().collect())
            },
        }
    }
}

pub(super) fn executable_schema_and_bindings(
    frontend_schema: &SchemaRef,
    required: &RequiredFields,
) -> Result<(SchemaRef, Vec<FieldBinding>), CapabilityReason> {
    let indices = required.requested_indices(frontend_schema.fields().len())?;
    let mut name_counts = BTreeMap::<&str, usize>::new();
    for index in &indices {
        let field = frontend_schema.fields().get(*index).ok_or(CapabilityReason::UnsupportedExpression)?;
        *name_counts.entry(field.name().as_str()).or_default() += 1;
    }

    let frontend_names = indices
        .iter()
        .map(|index| {
            frontend_schema
                .fields()
                .get(*index)
                .map(|field| field.name().clone())
                .ok_or(CapabilityReason::UnsupportedExpression)
        })
        .collect::<Result<BTreeSet<_>, _>>()?;

    let mut used_native_names = BTreeSet::new();
    let mut fields = Vec::with_capacity(indices.len());
    let mut bindings = Vec::with_capacity(indices.len());
    for frontend_index in indices {
        let field = frontend_schema.fields().get(frontend_index).ok_or(CapabilityReason::UnsupportedExpression)?;
        let frontend_name = field.name().clone();
        let native_name =
            if name_counts[frontend_name.as_str()] == 1 && !used_native_names.contains(frontend_name.as_str()) {
                frontend_name.clone()
            } else {
                unique_synthetic_field_name(frontend_index, frontend_name.as_str(), &frontend_names, &used_native_names)
            };
        used_native_names.insert(native_name.clone());
        fields.push(rename_field(field.as_ref(), native_name.as_str()));
        bindings.push(FieldBinding { frontend_index, frontend_name, native_name });
    }
    // Retain Arrow schema metadata when the native projection is an identity
    // over every frontend field. In-memory sources carry RecordBatches whose
    // schemas include that metadata, and the native source contract requires
    // those schemas to remain identical. Renamed or reduced projections get a
    // deliberately new executable schema instead.
    let preserves_frontend_schema = fields.len() == frontend_schema.fields().len()
        && bindings
            .iter()
            .enumerate()
            .all(|(index, binding)| binding.frontend_index == index && binding.frontend_name == binding.native_name);
    let schema = if preserves_frontend_schema { Arc::clone(frontend_schema) } else { Arc::new(Schema::new(fields)) };
    Ok((schema, bindings))
}

pub(super) fn rename_field(field: &Field, name: &str) -> Field {
    field.clone().with_name(name)
}

fn unique_synthetic_field_name(
    ordinal: usize,
    frontend_name: &str,
    frontend_names: &BTreeSet<String>,
    used_native_names: &BTreeSet<String>,
) -> String {
    let sanitized = sanitize_field_name(frontend_name);
    let base = format!("__algeon_native_field_{ordinal}_{sanitized}");
    unique_name_with_suffix(base.as_str(), |candidate| {
        !frontend_names.contains(candidate) && !used_native_names.contains(candidate)
    })
}