lutra-compiler 0.5.1

mod numbers;
mod repr;
mod repr_duckdb;
mod repr_pg;
mod types;

use std::collections::HashMap;

use lutra_bin::ir;
use lutra_sql as sa;

use crate::sql::utils::{Node, RelCols};
use crate::sql::{COL_ARRAY_INDEX, COL_VALUE, Dialect};
use crate::sql::{cr, utils};
use crate::utils::NameGenerator;

pub fn compile(
    rel: cr::Expr,
    types: HashMap<&ir::Path, &ir::Ty>,
    dialect: super::Dialect,
) -> sa::Query {
    let mut ctx = Context::new(types, dialect);

    let rel_ty = ctx.get_ty_mat(&rel.ty);
    let node = ctx.compile_rel(&rel);

    // convert to the repr used for query results
    let node = match dialect {
        super::Dialect::Postgres => {
            // PostgreSQL uses a repr very similar to "query repr", but:
            // - without array index column,
            // - ordered by array index column

            if rel_ty.kind.is_array() {
                let mut select = ctx.node_into_select(node, rel_ty);

                // drop index column
                let index = select.projection.remove(0);

                // apply the order
                let mut query = utils::query_select(select);
                query.order_by = utils::order_by_one(index.expr);
                Node::Query(query)
            } else {
                node
            }
        }

        super::Dialect::DuckDB => {
            // DuckDB uses "duckdb repr", so it can be converted to arrow
            ctx.duck_export(node, rel_ty, false)
        }
    };

    ctx.node_into_query(node, &rel.ty)
}

pub(super) struct Context<'a> {
    types: HashMap<&'a ir::Path, &'a ir::Ty>,

    pub(super) rel_name_gen: NameGenerator,

    rel_vars: HashMap<usize, RelRef>,

    dialect: super::Dialect,
}

struct RelRef {
    /// Name that can be used to refer to the relation
    name: String,

    /// True iff reference requires a FROM clause
    is_cte: bool,
}

impl<'a> Context<'a> {
    pub(super) fn new(types: HashMap<&'a ir::Path, &'a ir::Ty>, dialect: super::Dialect) -> Self {
        Self {
            types,
            rel_name_gen: NameGenerator::new("r"),
            rel_vars: Default::default(),
            dialect,
        }
    }

    pub(super) fn get_ty_mat(&self, ty: &'a ir::Ty) -> &'a ir::Ty {
        match &ty.kind {
            ir::TyKind::Ident(path) => self.types.get(path).unwrap(),
            _ => ty,
        }
    }

    pub(super) fn dialect(&self) -> super::Dialect {
        self.dialect
    }

    fn find_rel_var(&self, id: usize) -> &RelRef {
        match self.rel_vars.get(&id) {
            Some(x) => x,
            None => panic!("cannot find scope id: {id}"),
        }
    }

    fn register_rel_var(&mut self, input: &cr::BoundExpr, name: String) {
        let r = RelRef {
            name,
            is_cte: false,
        };
        self.rel_vars.insert(input.id, r);
    }

    fn register_cte(&mut self, input: &cr::BoundExpr, name: String) {
        let r = RelRef { name, is_cte: true };
        self.rel_vars.insert(input.id, r);
    }

    fn unregister_rel_var(&mut self, input: &cr::BoundExpr) {
        self.rel_vars.remove(&input.id);
    }

    #[tracing::instrument(name = "r", skip_all)]
    fn compile_rel(&mut self, rel: &cr::Expr) -> Node {
        match &rel.kind {
            cr::ExprKind::From(from) => self.compile_from(from, &rel.ty),

            cr::ExprKind::Join(left_in, right_in, condition) => {
                tracing::debug!("Join");

                let left = self.compile_rel(&left_in.rel);
                let (left_name, left_rels) = self.node_into_rel_var(left, &left_in.rel.ty);

                let right = self.compile_rel(&right_in.rel);
                let (right_name, right_rels) = self.node_into_rel_var(right, &right_in.rel.ty);

                let mut select = utils::select_empty();
                select.from.extend(left_rels);
                select.from.extend(right_rels);
                select.projection.extend(
                    self.projection(
                        &rel.ty,
                        Iterator::chain(
                            self.rel_cols(&left_in.rel.ty)
                                .map(|col| utils::identifier(Some(&left_name), col)),
                            self.rel_cols(&right_in.rel.ty)
                                .map(|col| utils::identifier(Some(&right_name), col)),
                        ),
                    ),
                );

                if let Some(condition) = condition {
                    self.register_rel_var(left_in, left_name);
                    self.register_rel_var(right_in, right_name);

                    select.selection = Some(self.compile_column(condition));

                    self.unregister_rel_var(left_in);
                    self.unregister_rel_var(right_in);
                }

                Node::Select(select)
            }

            cr::ExprKind::BindCorrelated(bound_in, main_in) => {
                tracing::debug!("BindCorrelated");

                // compile bound
                let bound = self.compile_rel(&bound_in.rel);
                let (bound_name, bound_rels) = self.node_into_rel_var(bound, &bound_in.rel.ty);

                // compile correlated
                self.register_rel_var(bound_in, bound_name);
                let main = self.compile_rel(main_in);
                self.unregister_rel_var(bound_in);

                let mut main = self.node_into_select(main, &main_in.ty);

                main.from = std::iter::Iterator::chain(
                    bound_rels.into_iter(),
                    main.from.into_iter().map(utils::lateral),
                )
                .collect();

                Node::Select(main)
            }

            cr::ExprKind::Transform(input_in, transform) => {
                tracing::debug!("Transform::{}", transform.as_ref());

                let mut input = self.compile_rel(&input_in.rel);

                let needs_input_rel = matches!(
                    transform,
                    cr::Transform::Aggregate(_)
                        | cr::Transform::Where(_)
                        | cr::Transform::Reindex(_)
                        | cr::Transform::Limit(..)
                        | cr::Transform::Group { .. }
                        | cr::Transform::Insert(_)
                );
                if needs_input_rel {
                    let (input_name, input_rel) = self.node_into_rel_var(input, &input_in.rel.ty);
                    input = input_rel
                        .map(Node::Rel)
                        .unwrap_or_else(|| Node::RelVar(input_name.clone()));

                    self.register_rel_var(input_in, input_name);
                }

                let r = self.compile_rel_transform(input, transform, &input_in.rel.ty, &rel.ty);

                if needs_input_rel {
                    self.unregister_rel_var(input_in);
                }

                r
            }

            cr::ExprKind::Bind(val_in, main_in) => {
                tracing::debug!("Bind");

                let name = self.rel_name_gen.next();

                let val = self.compile_rel(&val_in.rel);

                self.register_cte(val_in, name.clone());
                let main = self.compile_rel(main_in);
                self.unregister_rel_var(val_in);

                let val = self.node_into_query(val, &val_in.rel.ty);
                let mut main = self.node_into_query(main, &main_in.ty);

                let with = main.with.get_or_insert_with(utils::with);
                with.cte_tables.insert(0, utils::cte(name, val));

                Node::Query(main)
            }

            cr::ExprKind::Union(parts) => {
                tracing::debug!("Union");

                let mut res = None;
                for part_in in parts {
                    let part = self.compile_rel(part_in);

                    let part = self.node_into_query(part, &part_in.ty);

                    if let Some(r) = res {
                        res = Some(utils::union(r, utils::query_into_set_expr(part)))
                    } else {
                        res = Some(utils::query_into_set_expr(part))
                    }
                }

                let query =
                    utils::query_new(res.unwrap_or_else(|| self.construct_empty_rel(&rel.ty)));

                Node::Query(query)
            }

            cr::ExprKind::Iteration(initial_in, step_in) => {
                tracing::debug!("Iteration");

                let initial = self.compile_rel(initial_in);

                let re_name = self.rel_name_gen.next();

                self.register_cte(step_in, re_name.clone());
                let step = self.compile_rel(&step_in.rel);
                self.unregister_rel_var(step_in);

                let cte_query = utils::union(
                    utils::query_into_set_expr(self.node_into_query(initial, &initial_in.ty)),
                    utils::query_into_set_expr(self.node_into_query(step, &step_in.rel.ty)),
                );

                let mut main = utils::select_empty();
                main.projection = self.projection_noop(None, &rel.ty);
                main.from =
                    vec![sa::RelExpr::Table(utils::new_object_name([re_name.clone()])).unnamed()];

                let mut main = utils::query_select(main);
                main.with = Some(utils::with());

                let with = main.with.as_mut().unwrap();
                with.recursive = true;
                with.cte_tables
                    .push(utils::cte(re_name, utils::query_new(cte_query)));

                Node::Query(main)
            }

            cr::ExprKind::Update { table, updates } => {
                tracing::debug!("Update");

                let updates_ty = &updates.ty;
                let row_ty = updates.ty.kind.as_array().unwrap();

                let updates = self.compile_rel(updates);
                let updates = self.native_export(updates, updates_ty, true);
                let updates = self.node_into_rel(updates, updates_ty);
                let updates_var = utils::get_rel_alias(&updates).unwrap();

                // assignments: col1 = updates.col1, col2 = updates.col2
                let assignments = (self.native_cols(row_ty).into_iter())
                    .map(|(col, _ty)| sa::Assignment {
                        target: sa::AssignmentTarget::ColumnName(utils::new_object_name([&col])),
                        value: utils::identifier(Some(updates_var), col),
                    })
                    .collect();

                // where
                let selection = Some(utils::new_bin_op(
                    "=",
                    [
                        utils::identifier(Some(updates_var), "index"),
                        utils::identifier(Some(table), self.row_id()),
                    ],
                ));

                let table = utils::new_object_name([table]);
                let update = sa::SetExpr::Update(sa::Update {
                    table,
                    alias: None,
                    assignments,
                    from: vec![updates],
                    selection,
                    returning: None,
                    limit: None,
                });

                Node::Query(utils::query_new(update))
            }
        }
    }

    fn compile_from(&mut self, from: &cr::From, ty: &ir::Ty) -> Node {
        tracing::debug!("From::{}", from.as_ref());
        match from {
            cr::From::Row(row) => Node::Columns {
                exprs: self.compile_columns(row),
                rels: vec![],
            },

            cr::From::RelRef(scope_id) => {
                let rel_ref = self.find_rel_var(*scope_id);
                let rel_name = rel_ref.name.clone();

                if rel_ref.is_cte {
                    let alias = self.rel_name_gen.next();
                    let rel = sa::RelExpr::Table(utils::new_object_name([rel_name]))
                        .alias(utils::new_ident(alias));
                    Node::Rel(rel)
                } else {
                    Node::RelVar(rel_name)
                }
            }

            cr::From::Table { name, use_row_id } => {
                let node = Node::Rel(sa::RelExpr::Table(translate_table_ident(name)).unnamed());

                let mut node = self.native_import(node, ty);

                if *use_row_id {
                    // special case: use row id instead of ROW_NUMBER for index
                    let Node::Select(select) = &mut node else {
                        unreachable!()
                    };
                    select.projection[0].expr = utils::identifier(None::<&str>, self.row_id());
                }

                node
            }

            cr::From::Null => Node::from(self.null(ty)),
            cr::From::Literal(literal) => Node::from(self.compile_literal(literal, ty)),
            cr::From::FuncCall(func_name, args_in) => {
                self.compile_func_call(func_name, args_in, ty)
            }
            cr::From::Param(param_index) => {
                Node::Source(format!("${}::{}", param_index + 1, self.ty_name(ty)))
            }

            cr::From::Deserialize(input_in) => {
                let input = self.compile_rel(input_in);
                self.deserialize(input, &input_in.ty, ty)
            }

            cr::From::Serialize(expr_in) => {
                let expr = self.compile_rel(expr_in);
                self.serialize(expr, &expr_in.ty)
            }

            cr::From::Case(cases) => {
                let mut sql_cases = Vec::new();
                for (condition_in, result_in) in cases.iter().take(cases.len() - 1) {
                    let condition = self.compile_column(condition_in);

                    let result = self.compile_column(result_in);

                    sql_cases.push(sa::CaseWhen { condition, result });
                }

                // else
                let (_, else_value_in) = cases.last().unwrap();
                let else_value = self.compile_column(else_value_in);

                Node::from(sa::Expr::Case {
                    operand: None,
                    cases: sql_cases,
                    else_result: Some(Box::new(else_value)),
                })
            }

            cr::From::SQLSource(source) => {
                let node = Node::Source(source.clone());

                self.native_import(node, ty)
            }
        }
    }

    fn compile_rel_transform(
        &mut self,
        input: Node,
        transform: &cr::Transform,
        input_ty: &ir::Ty,
        output_ty: &ir::Ty,
    ) -> Node {
        let r = match transform {
            cr::Transform::ProjectPick(cols) => {
                let (mut columns, rels) = self.node_into_columns_and_rels(input, input_ty);

                utils::pick_by_position(&mut columns, cols);

                // apply new column names
                Node::Columns {
                    exprs: columns,
                    rels,
                }
            }

            cr::Transform::ProjectDiscard(cols) => {
                let (mut columns, rels) = self.node_into_columns_and_rels(input, input_ty);

                utils::drop_by_position(&mut columns, cols);

                // apply new column names
                Node::Columns {
                    exprs: columns,
                    rels,
                }
            }

            cr::Transform::Aggregate(columns) => {
                let input_rel = self.node_into_rel(input, input_ty);

                let (exprs, r) = self.compile_columns_scoped(columns);
                let mut rels = Vec::with_capacity(r.len() + 1);
                rels.insert(0, input_rel);
                rels.extend(r.into_iter().map(utils::lateral));
                Node::Columns { exprs, rels }
            }

            cr::Transform::Where(cond_in) => {
                // wrap into a new query
                let mut select = self.node_into_select(input, input_ty);

                let cond = self.compile_column(cond_in);
                utils::set_or_bin_op(&mut select.selection, "AND", cond);
                Node::Select(select)
            }
            cr::Transform::Limit(limit_in, order_by) => {
                let mut query = self.node_into_query(input, input_ty);

                query.limit = Some(utils::number(limit_in.to_string()));

                let order_by = self.compile_column(order_by);
                query.order_by = utils::order_by_one(order_by);

                Node::Query(query)
            }
            cr::Transform::Reindex(keys) => {
                // wrap into a new query
                let rel = self.node_into_rel(input, input_ty);
                let mut select = self.rel_into_select(rel, output_ty);

                // overwrite array index
                let keys = self.compile_columns(keys);

                select.projection[0] = sa::SelectItem {
                    expr: sa::Expr::IndexBy(keys),
                    alias: Some(utils::new_ident(COL_ARRAY_INDEX)),
                };
                Node::Select(select)
            }
            cr::Transform::Order => {
                let mut query = self.node_into_query(input, input_ty);

                // overwrite ORDER BY
                query.order_by =
                    utils::order_by_one(utils::identifier(None::<&str>, COL_ARRAY_INDEX));
                Node::Query(query)
            }

            cr::Transform::Group {
                key,
                values: values_in,
            } => {
                // wrap into a new query
                let rel = self.node_into_rel(input, input_ty);
                let mut select = self.rel_into_select(rel, output_ty);

                let group_by = self.compile_columns(core::slice::from_ref(key));

                let mut projection = vec![
                    sa::Expr::IndexBy(vec![]), // index
                ];

                // values
                let (values, val_rels) = self.compile_columns_scoped(values_in);

                projection.extend(values);
                select.from.extend(val_rels.into_iter().map(utils::lateral));

                select.group_by = group_by;
                select.projection = self.projection(output_ty, projection);

                Node::Select(select)
            }

            cr::Transform::Insert(table_ident) => {
                let source = self.native_export(input, input_ty, false);
                let source = self.node_into_query(source, input_ty);

                let table = translate_table_ident(table_ident);

                let columns = self
                    .native_cols(input_ty)
                    .into_iter()
                    .map(|(f_name, _f_ty)| f_name)
                    .map(utils::new_ident)
                    .collect();

                let insert = sa::Insert {
                    source: Box::new(source),
                    table,
                    columns,
                };
                let query = utils::query_new(sa::SetExpr::Insert(insert));
                Node::Query(query)
            }
        };
        tracing::trace!("-> {r:?}");
        r
    }

    fn compile_columns_scoped(
        &mut self,
        columns: &[cr::Expr],
    ) -> (Vec<sa::Expr>, Vec<sa::RelNamed>) {
        let mut exprs = Vec::with_capacity(columns.len());
        let mut relations = Vec::new();
        for col in columns {
            let node = self.compile_rel(col);
            let (expr, rels) = self.node_into_columns_and_rels(node, &col.ty);
            exprs.extend(expr);
            relations.extend(rels);
        }
        (exprs, relations)
    }

    fn compile_columns(&mut self, columns: &[cr::Expr]) -> Vec<sa::Expr> {
        let mut column_exprs = Vec::with_capacity(columns.len());
        for col in columns {
            let value = self.compile_rel(col);
            column_exprs.extend(self.node_into_columns(value, &col.ty));
        }
        column_exprs
    }

    /// Compile an expression as relation and then fit the result into a column.
    /// In the worst case, this creates a subquery.
    fn compile_column(&mut self, expr: &cr::Expr) -> sa::Expr {
        let rel = self.compile_rel(expr);
        self.node_into_column(rel, &expr.ty)
    }

    fn compile_func_call(&mut self, id: &str, args_in: &[cr::Expr], ty: &ir::Ty) -> Node {
        let args = self.compile_columns(args_in);
        let expr = match id {
            "std::mul" => utils::new_bin_op("*", args),
            "std::div" => match self.dialect {
                Dialect::Postgres => utils::new_bin_op("/", args),
                Dialect::DuckDB => match ty.kind.as_primitive().unwrap() {
                    ir::TyPrimitive::float32 | ir::TyPrimitive::float64 => {
                        utils::new_bin_op("/", args)
                    }
                    _ => utils::new_bin_op("//", args),
                },
            },
            "std::mod" => match ty.kind.as_primitive().unwrap() {
                ir::TyPrimitive::float32 | ir::TyPrimitive::float64 => {
                    let [l, r] = unpack_args(args);
                    sa::Expr::Source(format!("MOD({l}::numeric, {r}::numeric)::float8"))
                }
                _ => utils::new_bin_op("%", args),
            },
            "std::add" => utils::new_bin_op("+", args),
            "std::sub" => utils::new_bin_op("-", args),
            "std::neg" => utils::new_un_op("-", args),

            "std::cmp" => {
                let [a, b] = unpack_args(args);
                sa::Expr::Source(format!(
                    "(SELECT CASE WHEN a < b THEN 0 WHEN a > b THEN 2 ELSE 1 END::int2 FROM (VALUES ({a}, {b})) t(a,b))"
                ))
            }
            "std::eq" => utils::new_bin_op("=", args),
            "std::lt" => utils::new_bin_op("<", args),
            "std::lte" => utils::new_bin_op("<=", args),

            "std::and" => utils::new_bin_op("AND", args),
            "std::or" => utils::new_bin_op("OR", args),
            "std::not" => utils::new_un_op("NOT", args),

            "std::sequence" => {
                let [start, end] = unpack_args(args);

                // generate_series
                let generate_series = sa::RelExpr::function(
                    utils::new_ident("generate_series"),
                    vec![start.clone(), sa::Expr::Source(format!("({end} - 1)"))],
                );

                // select index, value from generate_series()
                let mut select = utils::select_empty();
                select.from.push(sa::RelNamed {
                    lateral: false,
                    expr: generate_series,
                    alias: Some(sa::TableAlias {
                        name: sa::Ident::new("t"),
                        columns: vec![sa::Ident::new(COL_VALUE)],
                    }),
                });
                let out_cast_ty = self.ty_name(&args_in[0].ty);
                select.projection = vec![
                    sa::SelectItem {
                        expr: sa::Expr::Source(format!("(t.{COL_VALUE} - {start})::int4")),
                        alias: Some(COL_ARRAY_INDEX.into()),
                    },
                    sa::SelectItem {
                        expr: sa::Expr::Source(format!("t.{COL_VALUE}::{out_cast_ty}")),
                        alias: Some(COL_VALUE.into()),
                    },
                ];
                return Node::Select(select);
            }

            "std::min" => {
                let ty_item = self.get_ty_mat(&args_in[0].ty);

                let is_bool = ty_item
                    .kind
                    .as_primitive()
                    .is_some_and(|p| *p == ir::TyPrimitive::bool);

                if is_bool {
                    // special case: bool
                    utils::new_func_call("BOOL_AND", args)
                } else {
                    // base case
                    utils::new_func_call("MIN", args)
                }
            }
            "std::max" => {
                let ty_item = self.get_ty_mat(&args_in[0].ty);

                let is_bool = ty_item
                    .kind
                    .as_primitive()
                    .is_some_and(|p| *p == ir::TyPrimitive::bool);

                if is_bool {
                    // special case: bool
                    utils::new_func_call("BOOL_OR", args)
                } else {
                    // base case
                    utils::new_func_call("MAX", args)
                }
            }
            "std::sum" => {
                let [arg] = unpack_args(args);
                let ty = self.ty_name(ty);
                sa::Expr::Source(format!("COALESCE(SUM({arg}), 0)::{ty}"))
            }
            "std::mean" => {
                let [arg] = unpack_args(args);
                sa::Expr::Source(
                    if matches!(
                        args_in[0].ty.kind,
                        ir::TyKind::Primitive(ir::TyPrimitive::float32)
                    ) {
                        // case for: float4, int2, int4, int8
                        // PostgreSQL does not have a specialized float4 avg, but just uses float8
                        // This yields different results that it should.
                        // It also has exact precision for integers via numeric type, which is
                        // more precise, but slower than using floats for division.
                        // So we implement our own AVG()
                        format!("SUM({arg})/COUNT({arg})")
                    } else {
                        // case for: float8
                        format!("COALESCE(AVG({arg})::float8, 'NaN')")
                    },
                )
            }
            "std::count" => {
                let [arg] = unpack_args(args);

                sa::Expr::Source(if let sa::Expr::CompoundIdentifier(parts) = arg {
                    let rvar = &parts[0];
                    format!("COUNT({rvar}.*)")
                } else {
                    // We need to reference arg, so we get aggregation of correct cardinality.
                    // But we cannot just use COUNT({arg}) because that will not count NULL values
                    // (which are produced by enum {None, Some: T} type).
                    // So we use IS NULL, which should be performant, but also always emit non-null values.
                    format!("COUNT({arg} IS NULL)")
                })
            }
            "std::any" => sa::Expr::Source(format!(
                "COALESCE({}, FALSE)",
                utils::new_func_call("BOOL_OR", args)
            )),
            "std::all" => sa::Expr::Source(format!(
                "COALESCE({}, TRUE)",
                utils::new_func_call("BOOL_AND", args)
            )),

            "std::lead" => {
                let [arg, offset] = unpack_args(args);

                let item_ty = self.get_ty_mat(ty).kind.as_array().unwrap();
                let filler = self.default_value(item_ty);

                sa::Expr::Source(format!(
                    "COALESCE(LEAD({arg}, {offset}::int4) OVER (ORDER BY {COL_ARRAY_INDEX}), {filler})"
                ))
            }
            "std::lag" => {
                let [arg, offset] = unpack_args(args);

                let item_ty = self.get_ty_mat(ty).kind.as_array().unwrap();
                let filler = self.default_value(item_ty);

                sa::Expr::Source(format!(
                    "COALESCE(LAG({arg}, {offset}::int4) OVER (ORDER BY {COL_ARRAY_INDEX}), {filler})"
                ))
            }
            "std::rolling_mean" => {
                let [array, trailing, leading] = unpack_args(args);

                sa::Expr::Source(format!(
                    "(AVG({array}) OVER (ORDER BY {COL_ARRAY_INDEX} ROWS BETWEEN {trailing} PRECEDING AND {leading} FOLLOWING))::float8"
                ))
            }
            "std::rank" => {
                let [array] = unpack_args(args);
                sa::Expr::Source(format!("(RANK() OVER (ORDER BY {array}))::int4"))
            }
            "std::rank_dense" => {
                let [array] = unpack_args(args);
                sa::Expr::Source(format!("(DENSE_RANK() OVER (ORDER BY {array}))::int4"))
            }
            "std::rank_percentile" => {
                let [array] = unpack_args(args);
                sa::Expr::Source(format!("PERCENT_RANK() OVER (ORDER BY {array})"))
            }
            "std::cume_dist" => {
                let [array] = unpack_args(args);
                sa::Expr::Source(format!("CUME_DIST() OVER (ORDER BY {array})"))
            }

            "std::text::length" => {
                let [text] = unpack_args(args);
                sa::Expr::Source(format!(
                    "LENGTH({text})::{}",
                    self.ty_name(&ir::Ty::new(ir::TyPrimitive::uint32))
                ))
            }
            "std::text::from_ascii" => {
                let [ascii] = unpack_args(args);
                sa::Expr::Source(format!("CHR({ascii})"))
            }
            "std::text::concat" => utils::new_bin_op("||", args),
            "std::text::join" => {
                let [parts, sep] = unpack_args(args);
                sa::Expr::Source(format!("COALESCE(STRING_AGG({parts}, {sep}), '')"))
            }
            "std::text::split" => {
                let [text, sep] = unpack_args(args);

                match self.dialect {
                    Dialect::Postgres => {
                        // PostgreSQL: use string_to_table function
                        let mut select = utils::select_empty();

                        select.from.push(sa::RelNamed {
                            lateral: false,
                            alias: Some(sa::TableAlias {
                                name: sa::Ident::new("t"),
                                columns: vec![sa::Ident::new(COL_VALUE)],
                            }),
                            expr: sa::RelExpr::function(
                                utils::new_ident("string_to_table"),
                                vec![text, sep],
                            ),
                        });
                        select.projection = vec![
                            sa::SelectItem {
                                expr: sa::Expr::IndexBy(vec![]),
                                alias: Some(COL_ARRAY_INDEX.into()),
                            },
                            sa::SelectItem::unnamed(utils::identifier(Some("t"), COL_VALUE)),
                        ];
                        return Node::Select(select);
                    }
                    Dialect::DuckDB => {
                        // DuckDB: string_split + unnest

                        let split = sa::Expr::Source(format!("split({text}, {sep})"));

                        // convert into a relation via unnest
                        let unnest = sa::RelExpr::Function {
                            name: utils::new_object_name(["unnest"]),
                            args: vec![split],
                            ordinality: true,
                        };

                        let mut select = utils::select_empty();
                        select.from.push(unnest.alias_cols(
                            utils::new_ident("t"),
                            vec![utils::new_ident(COL_VALUE), utils::new_ident("idx")],
                        ));
                        select.projection = vec![
                            sa::SelectItem {
                                expr: sa::Expr::Source("t.idx - 1".into()),
                                alias: Some(COL_ARRAY_INDEX.into()),
                            },
                            sa::SelectItem::unnamed(utils::identifier(Some("t"), COL_VALUE)),
                        ];
                        return Node::Select(select);
                    }
                }
            }
            "std::text::starts_with" => utils::new_func_call("STARTS_WITH", args),
            "std::text::ends_with" => {
                let [text, suffix] = unpack_args(args);
                sa::Expr::Source(format!("({text} LIKE '%' || {suffix})"))
            }
            "std::text::contains" => {
                let [text, pattern] = unpack_args(args);
                sa::Expr::Source(format!("({text} LIKE '%' || {pattern} || '%')"))
            }

            "std::math::abs" => {
                let [text] = unpack_args(args);
                sa::Expr::Source(format!("ABS({text})"))
            }
            "std::math::pow" => {
                let [operand, exponent] = unpack_args(args);
                sa::Expr::Source(format!("POW({operand}, {exponent})"))
            }

            "greatest" => {
                let mut args = args.into_iter();
                sa::Expr::Source(format!(
                    "GREATEST({}, {})::int8",
                    args.next().unwrap(),
                    args.next().unwrap(),
                ))
            }

            "is_null" => {
                let [arg] = unpack_args(args);
                sa::Expr::Source(format!("{arg} IS NULL"))
            }
            "is_not_null" => {
                let [arg] = unpack_args(args);
                sa::Expr::Source(format!("{arg} IS NOT NULL"))
            }

            "std::to_int8" | "std::to_int16" | "std::to_int32" | "std::to_int64"
            | "std::to_uint8" | "std::to_uint16" | "std::to_uint32" | "std::to_uint64" => {
                let [arg] = unpack_args(args);
                self.compile_to_int(arg, &args_in[0].ty, ty)
            }

            "std::to_float32" | "std::to_float64" => {
                let [arg] = unpack_args(args);
                let ty = self.ty_name(ty);
                sa::Expr::Source(format!("({arg})::{ty}"))
            }

            "std::to_text" => {
                let [arg] = unpack_args(args);
                let ty = self.ty_name(ty);

                let mut r = sa::Expr::Source(format!("({arg})::{ty}"));

                if let Dialect::DuckDB = self.dialect {
                    // DuckDB includes .0 for whole number floats -> remove it

                    let arg_ty = self.get_ty_mat(&args_in[0].ty);
                    let is_float = matches!(
                        arg_ty.kind,
                        ir::TyKind::Primitive(ir::TyPrimitive::float32 | ir::TyPrimitive::float64)
                    );
                    if is_float {
                        r = sa::Expr::Source(format!("RTRIM(RTRIM({r}, '0'), '.')"));
                    }
                }

                r
            }

            "std::fs::read_parquet" => match self.dialect {
                Dialect::DuckDB => {
                    // Only DuckDB supports read_parquet
                    let [file_name] = unpack_args(args);
                    let node = Node::Rel(sa::RelNamed::unnamed(sa::RelExpr::function(
                        utils::new_ident("read_parquet"),
                        vec![file_name],
                    )));
                    return self.native_import(node, ty);
                }
                Dialect::Postgres => {
                    panic!("read_parquet is not supported on PostgreSQL")
                }
            },
            "std::fs::write_parquet" => match self.dialect {
                Dialect::DuckDB => {
                    // Only DuckDB supports write_parquet
                    let [data, file_name] = unpack_args(args);

                    // For write_parquet, we need to export the data and then write it
                    let data = self.native_export(Node::from(data), &args_in[0].ty, false);
                    let data = self.node_into_query(data, &args_in[0].ty);

                    return Node::Query(utils::query_new(sa::SetExpr::Copy(Box::new(sa::Copy {
                        source: sa::SetExpr::Query(Box::new(data)),
                        target: file_name,
                        options: "FORMAT parquet, COMPRESSION zstd".into(),
                    }))));
                }
                Dialect::Postgres => {
                    panic!("write_parquet is not supported on PostgreSQL")
                }
            },

            "std::date::to_timestamp" => {
                let [date, tz] = unpack_args(args);

                let timestamp =
                    format!("('1970-01-01'::date + {date})::timestamp AT TIME ZONE {tz}");

                sa::Expr::Source(format!("(EXTRACT(EPOCH FROM {timestamp}) * 1000000)::int8"))
            }
            "std::date::to_year_month_day" => {
                let [date] = unpack_args(args);
                let f = self.get_ty_mat(ty).kind.as_tuple().unwrap();

                let mut input_select = utils::select_empty();
                input_select.projection = vec![sa::SelectItem {
                    expr: sa::Expr::Source(format!("'1970-01-01'::date + {date}")),
                    alias: Some("x".into()),
                }];
                let input_query = utils::query_select(input_select);

                let mut select = utils::select_empty();
                select
                    .from
                    .push(sa::RelExpr::subquery(input_query).alias(utils::new_ident("t")));

                let values = [
                    format!("date_part('YEAR', t.x)::{}", self.ty_name(&f[0].ty)),
                    format!("date_part('MONTH', t.x)::{}", self.ty_name(&f[1].ty)),
                    format!("date_part('DAY', t.x)::{}", self.ty_name(&f[2].ty)),
                ]
                .into_iter()
                .map(sa::Expr::Source);
                select.projection = self.projection(ty, values);
                return Node::Select(select);
            }
            "std::timestamp::to_date" => {
                let [timestamp, tz] = unpack_args(args);

                let local_date = match self.dialect {
                    Dialect::Postgres => {
                        let timestamp = format!("to_timestamp({timestamp}::float8/1000000.0)");
                        format!("({timestamp} AT TIME ZONE {tz})::date")
                    }
                    Dialect::DuckDB => {
                        let timestamp = format!("make_timestamp({timestamp})");
                        format!("(({timestamp} AT TIME ZONE 'UTC') AT TIME ZONE {tz})::date")
                    }
                };

                sa::Expr::Source(format!("({local_date} - '1970-01-01'::date)::int4"))
            }

            _ => todo!("sql impl for {id}"),
        };
        Node::from(expr)
    }

    fn null(&self, ty: &ir::Ty) -> sa::Expr {
        sa::Expr::Source(format!("NULL::{}", self.ty_name(ty)))
    }

    fn construct_empty_rel(&self, ty: &ir::Ty) -> sa::SetExpr {
        // construct a rel without rows
        let mut values = Vec::new();

        if ty.kind.is_array() {
            values.push(utils::number("0"));
        }
        let item_ty = match &ty.kind {
            ir::TyKind::Primitive(_) | ir::TyKind::Tuple(_) => ty,
            ir::TyKind::Array(item_ty) => item_ty,
            ir::TyKind::Enum(_) | ir::TyKind::Function(_) | ir::TyKind::Ident(_) => {
                todo!()
            }
        };
        match &self.get_ty_mat(item_ty).kind {
            ir::TyKind::Primitive(_) => {
                values.push(self.null(item_ty));
            }
            ir::TyKind::Tuple(ty_fields) => {
                values.extend(ty_fields.iter().map(|ty_field| self.null(&ty_field.ty)));
            }
            _ => todo!(),
        }

        let mut select = utils::select_empty();
        select.projection = self.projection(ty, values);
        select.selection = Some(utils::bool(false));
        sa::SetExpr::Select(Box::new(select))
    }
}

fn unpack_args<const N: usize>(args: impl IntoIterator<Item = sa::Expr>) -> [sa::Expr; N] {
    let mut r = Vec::with_capacity(N);
    let mut args = args.into_iter();
    for _ in 0..N {
        r.push(args.next().unwrap());
    }
    r.try_into().unwrap()
}

/// Converts a "lutra table identifier" into an SQL object name.
///
/// For example:
/// - `"invoices" -> invoices`
/// - `"mySchema/invoices" -> "mySchema".invoices`
/// - `"my_schema/hello/world" -> my_schema."hello/world"`
/// - `"public/my_dir/sub_dir/data.parquet" -> public."my_dir/sub_dir/data.parquet"`
///
/// Splits the name on first slash and uses preceding part (if it exists)
/// for the schema name, and following part as the table name.
/// Thus, table name can contain slashes.
///
/// This function will be specialized for different databases,
/// because current behavior targets PostgreSQL.
/// For example, BigQuery uses organizations, projects and datasets in table identifiers.
/// So when translating to BigQuery SQL we will translate:
/// - `"my_organization/my_project/my_dataset/my_table" -> my_organization.my_project.my_dataset.my_table`,
fn translate_table_ident(table_ident: &str) -> sa::ObjectName {
    if let Some((schema, table)) = table_ident.split_once('/') {
        utils::new_object_name([schema, table])
    } else {
        utils::new_object_name([table_ident])
    }
}