linkrs-query 0.1.0

Query engine, parser, and executor for Linkrs
Documentation
use std::collections::HashMap;

use linkrs_core::types::expr::ExpressionId;

use crate::planning::plan::factorization::{
    FGroupPos, FactorizationError, FactorizedSchema, SchemaUtils,
};

use crate::planning::plan::logical::logical_nodes::graph_ops::LogicalAssignNode;

pub(super) fn assign(
    n: &LogicalAssignNode,
    child_schemas: &[FactorizedSchema],
) -> Result<FactorizedSchema, FactorizationError> {
    let schema = child_schemas.first().cloned().unwrap_or_default();
    if schema.num_groups() == 0 {
        let mut out = FactorizedSchema::new();
        let g = out.create_flat_group(false);
        for (alias, expr) in &n.assignments {
            let eid = expr.id().clone();
            out.insert_to_group_and_scope_with_name(eid, Some(alias.clone()), g)?;
        }
        out.validate_at_most_one_unflat()?;
        return Ok(out);
    }
    let mut expr_store: HashMap<ExpressionId, linkrs_core::Expression> = HashMap::new();
    for (_, expr) in &n.assignments {
        if let Some(rhs) = expr.get_expression() {
            expr_store.insert(expr.id().clone(), rhs);
        }
    }
    let mut out = schema.clone();
    for (alias, expr) in &n.assignments {
        let alias_id = expr.id().clone();
        if out.is_expression_in_scope(&alias_id) {
            continue;
        }
        let mut analyzer =
            crate::optimizer::factorization::GroupDependencyAnalyzer::with_expr_store(
                &out,
                false,
                &expr_store,
            );
        analyzer.visit(&alias_id);
        let dependent = analyzer.dependent_groups().clone();
        let required_flat = analyzer.required_flat_groups().clone();
        for pos in required_flat.iter() {
            if let Some(g) = out.get_group(*pos) {
                if !g.is_flat() {
                    out.flatten_group(*pos)?;
                }
            }
        }
        let target = if dependent.is_empty() {
            out.groups()
                .iter()
                .enumerate()
                .find(|(_, g)| g.is_flat())
                .map(|(i, _)| i as FGroupPos)
                .unwrap_or_else(|| out.create_flat_group(false))
        } else if dependent.len() == 1 {
            *dependent.iter().next().unwrap()
        } else {
            let mut candidates: Vec<FGroupPos> = dependent
                .iter()
                .filter(|pos| {
                    out.get_group(**pos)
                        .map(|g| !g.is_flat() && !required_flat.contains(pos))
                        .unwrap_or(false)
                })
                .copied()
                .collect();
            candidates.sort_unstable();
            if candidates.is_empty() {
                SchemaUtils::get_leading_group_pos(&dependent, &out)?
            } else {
                candidates[0]
            }
        };
        out.insert_to_scope_with_name(alias_id.clone(), alias.clone(), target)?;
        if let Some(g) = out.get_group_mut(target) {
            if !g.contains(&alias_id) {
                g.insert_expression_with_name(alias_id.clone(), Some(alias.clone()))?;
            }
        }
    }
    out.validate_at_most_one_unflat()?;
    Ok(out)
}