///|
/// Index a frame's rows by composite join key: maps each non-null key tuple
/// to the ascending list of row indices that carry it. A row whose key has
/// any null cell is unmatchable (`join_row_key` returns `None`), so it is
/// never indexed. Shared by every `how` so the null-key skip and the
/// bucket-append convention live in one place; bucket order is ascending
/// because the row index only grows. Total.
fn index_by_join_key(
  key_scalars : Array[Array[@types.Scalar]],
  n : Int,
) -> Map[Array[KeyCell], Array[Int]] {
  let index : Map[Array[KeyCell], Array[Int]] = Map([])
  for i in 0..
        match index.get(k) {
          Some(rows) => rows.push(i)
          None => index[k] = [i]
        }
      None => ()
    }
  }
  index
}

///|
/// Equi-join `self` (the left frame) with `other` (the right frame) on the
/// key **expressions** the options carry — `JoinOptions::on(keys)`, or the
/// paired `JoinOptions::left_on(keys, right_on=keys)` — producing a new
/// `DataFrame`.
///
/// Each key is an arbitrary expression, evaluated over the whole frame
/// under the rules in `expr_eval.mbt`, exactly like a `sort` / `group_by`
/// key: an `on` key is applied to both frames (a bare `col("id")` joins on
/// an existing column, a derived key such as `col("ts") / lit_int(86400)`
/// joins on the computed value), while a sided pair evaluates
/// position-paired keys on the left and right frame respectively (for
/// differently-named or differently-derived keys). A key that reduces to a
/// single cell (a literal, or an aggregation) broadcasts over the frame.
/// Reading back what an options value carries is `on_keys()` /
/// `left_keys()` / `right_keys()`, each returning a copy.
///
/// Two rows match when every key holds an equal value, using the same
/// composite-`KeyCell`-tuple encoding `group_by` uses (see
/// `frame/row_key.mbt`; the tuple is structurally injective across key
/// columns). The one deliberate difference from `group_by`: a **null** key
/// cell matches **nothing** (`null != null`, the SQL / Polars default) —
/// such an unmatched row is dropped by `Inner` and kept (with the other
/// side's columns null) by `Left` / `Right` / `Outer`. A `Float` `NaN` key
/// is *not* null, so — as in `group_by`, and matching Polars' "NaN compares
/// equal" rule — all `NaN` keys match each other.
///
/// `how` selects which unmatched rows survive:
///   * `Inner` — only matched pairs.
///   * `Left` — matched pairs plus every unmatched **left** row (right
///     columns null).
///   * `Right` — matched pairs plus every unmatched **right** row (left
///     columns null); the mirror of `Left`.
///   * `Outer` — matched pairs plus every unmatched row from **both** sides.
///   * `Cross` — the keyless Cartesian product (see below).
///
/// Output shape:
///   * **columns** = the left columns (original order and names) followed
///     by the right frame's columns (original order). Coalescing applies
///     only to an `on` join whose every key is a bare `col(...)` (the only
///     shape where a key names the same column on both sides); `left_on` /
///     `right_on` and any derived key never coalesce. Whether an eligible
///     right key column is kept is then governed by `options.coalesce`
///     (`None` = auto by `how`, matching Polars: an inner / left / right join
///     coalesces, an outer join does not). When coalesced, the key appears
///     once at the left key's position, taking each row's value from
///     whichever side is present (the left on `Inner` / `Left`, the right
///     on `Right`, the present side per row on `Outer` — the two are equal
///     on a matched pair); the right key column is dropped. When not
///     coalesced, the right key column is kept too — it clashes with the
///     left key name, so it gains `options.suffix` (e.g. `id_right`) and is
///     null wherever its row had no match. Any other right column whose
///     name occurs in the left frame is likewise suffixed; the left column
///     keeps its name. (A derived key contributes no column of its own — it
///     only decides which rows pair — so the output is just the two frames'
///     columns, the right suffixed on a clash.)
///   * **rows** = left rows in their original order (each with its right
///     matches in ascending right-row order, then — for `Left` / `Outer` —
///     unmatched left rows in place with null right columns), followed for
///     `Outer` by the unmatched right rows in right-row order. A `Right`
///     join instead emits every right row in right-row order (each with its
///     left matches in ascending left-row order, else the right row alone
///     with null left columns). The order is fully determined by the input
///     order, so results are snapshot-stable.
///
/// `how = Cross` is the **Cartesian product** (every left row paired with
/// every right row); it takes no keys, ignores `coalesce`, and keeps all
/// columns of both frames (suffixing a right column that clashes with a
/// left name). It is the explicit form of what `group_by([])`'s
/// grand-total group is for aggregation.
///
/// The output height is the row plan's own length, never inferred from the
/// assembled columns, so a join between two column-less frames keeps its
/// rows: `2×0` cross `3×0` is the `6×0` frame, not `0×0`.
///
/// The result routes through `DataFrame::from_parts`, so it satisfies
/// `check_invariants()`.
///
/// Raises:
/// - `ColumnNotFound(name)` — a key expression references an absent column.
///   Reported on the first offending key, in key order, evaluating the left
///   frame's key before the right's.
/// - `TypeMismatch(detail)` — a key's left and right dtypes differ (so its
///   values could never compare equal), or a derived key's own dtypes don't
///   unify. Reported on the first such key, in key order.
/// - `InvalidOperation(detail)` — no keys for a non-`Cross` join (use
///   `how = Cross` for a Cartesian product), any keys on a `Cross` join
///   (which takes none), or a sided pair of unequal length. (Mixing shared
///   and sided keys is a fourth case the engine still checks, though the
///   three constructors make it unspellable from outside `frame`.)
/// - `DuplicateColumn(name)` — either a repeated shared key that names a
///   column twice (only an `on` join of bare `col` keys is deduplicated, so
///   only that form can reach this, like `group_by([col("id"), col("id")])`),
///   or two output columns still
///   colliding after suffixing (e.g. the left frame already has both
///   `value` and `value_right`, and the right contributes a non-key
///   `value`; surfaced by `DataFrame::from_parts`).
/// - `LengthMismatch` — a `lit_series(s)` key whose embedded series is
///   neither length 1 nor the frame's height (the evaluator's broadcast
///   rule, surfacing through key evaluation).
pub fn DataFrame::join(
  self : DataFrame,
  other : DataFrame,
  options : JoinOptions,
) -> DataFrame raise @types.DataError {
  let (left_rows, right_rows) = self.plan_join_rows(other, options)
  self.assemble_join_output(other, options, left_rows, right_rows)
}

///|
/// Phase 1 of `join`: build the output row plan as parallel optional-index
/// vectors — for output row `r`, `left_rows[r]` is its left row (`None` = a
/// right row kept with no left match) and `right_rows[r]` its right row
/// (`None` = a left row kept with no right match). At least one side is
/// always `Some`. A `Cross` join is the keyless Cartesian product; the key
/// joins hash one frame by key and probe it with the other. Split from the
/// column assembly (`assemble_join_output`) so row planning and output
/// materialisation each read on their own.
fn DataFrame::plan_join_rows(
  self : DataFrame,
  other : DataFrame,
  options : JoinOptions,
) -> (Array[Int?], Array[Int?]) raise @types.DataError {
  let left_rows : Array[Int?] = []
  let right_rows : Array[Int?] = []
  match options.how {
    Cross => {
      // A cross join takes no keys (Polars rejects keys with `how="cross"`).
      if !options.on.is_empty() ||
        !options.left_on.is_empty() ||
        !options.right_on.is_empty() {
        raise @types.DataError::InvalidOperation(
          "cross join takes no key columns",
        )
      }
      // Every left row paired with every right row, both in input order.
      for i in 0.. {
      // Mirror of `Left`: hash the LEFT frame by key, then probe it with
      // each right row in order. The row set is every right row — its left
      // matches (in ascending left-row order), else the right row alone with
      // null left columns. Output column order is unchanged (left then
      // right); only the surviving rows differ.
      let (left_key_scalars, right_key_scalars) = resolve_join_keys(
        self, other, options,
      )
      // Hash the left frame: composite key -> ascending matching left rows
      // (null-key rows unmatchable, so never indexed).
      let left_index = index_by_join_key(left_key_scalars, self.nrows())
      // Probe each right row in order. A null right key (`None`) and a key
      // with no left entry are the same case: no match, keep the right row
      // with null left columns.
      for j in 0.. {
          left_index.get(k)
        })
        match matches {
          Some(rows) =>
            for i in rows {
              left_rows.push(Some(i))
              right_rows.push(Some(j))
            }
          None => {
            left_rows.push(None)
            right_rows.push(Some(j))
          }
        }
      }
    }
    Inner | Left | Outer => {
      let (left_key_scalars, right_key_scalars) = resolve_join_keys(
        self, other, options,
      )
      // Hash the right frame: composite key -> ascending matching right rows
      // (null-key rows unmatchable, so never indexed).
      let right_index = index_by_join_key(right_key_scalars, other.nrows())
      // `Left` / `Outer` keep an unmatched left row (null right); `Inner`
      // drops it. `Outer` additionally appends the unmatched right rows.
      let keep_unmatched_left = match options.how {
        Left | Outer => true
        Inner | Right | Cross => false
      }
      let emit_unmatched_right = match options.how {
        Outer => true
        Inner | Left | Right | Cross => false
      }
      // Track which right rows were paired, so `Outer` can emit the rest.
      // Only `Outer` reads this, so `Inner` / `Left` skip both the buffer and
      // the per-match write below.
      let right_matched = if emit_unmatched_right {
        Array::make(other.nrows(), false)
      } else {
        []
      }
      // Probe each left row in order. A null left key (`None`) and a key
      // with no right entry are the same case: no match.
      for i in 0.. {
          right_index.get(k)
        })
        match matches {
          Some(rows) =>
            for j in rows {
              left_rows.push(Some(i))
              right_rows.push(Some(j))
              if emit_unmatched_right {
                right_matched[j] = true
              }
            }
          None =>
            if keep_unmatched_left {
              left_rows.push(Some(i))
              right_rows.push(None)
            }
        }
      }
      // Full outer: append every right row that never paired, in right-row
      // order, with null left columns.
      if emit_unmatched_right {
        for j in 0.. null`), right-column
/// suffixing on a name clash with the left frame, and per-column backend
/// re-convergence (`preserve_backend`). Split from the row planning
/// (`plan_join_rows`) so each phase reads on its own.
///
/// The height is the plan's — `left_rows.length()`, which every gather above
/// produces a column of — rather than the assembled columns'. The two agree
/// wherever a column exists, and where none does (both frames column-less)
/// only the plan carries the height.
fn DataFrame::assemble_join_output(
  self : DataFrame,
  other : DataFrame,
  options : JoinOptions,
  left_rows : Array[Int?],
  right_rows : Array[Int?],
) -> DataFrame raise @types.DataError {
  // Resolve key coalescing. Only an `on` join whose every key is a bare
  // `col(...)` can coalesce — that is the one shape where a key names the
  // same column on both sides (`coalesce_key_names` returns those names);
  // `left_on` / `right_on` and derived keys turn coalescing off. When
  // eligible, an explicit `coalesce` argument wins, else it is auto by `how`
  // (inner / left / right coalesce → key once; outer / cross do not),
  // matching Polars' `coalesce = None`. `on_set` is the set of names to
  // coalesce — empty whenever coalescing is off, so the assembly below
  // keeps every column.
  let coalesce_names = coalesce_key_names(options)
  let coalesce = if coalesce_names is Some(_) {
    options.coalesce.unwrap_or(
      match options.how {
        Inner | Left | Right => true
        Outer | Cross => false
      },
    )
  } else {
    false
  }
  let on_set : Map[String, Unit] = Map([])
  if coalesce_names is Some(names) {
    for k in names {
      on_set[k] = ()
    }
  }
  let left_names : Map[String, Unit] = Map([])
  for n in self.columns() {
    left_names[n] = ()
  }
  // Assemble the output columns. Left columns are gathered with `None ->
  // null` (a kept unmatched-right row has no left cell). A coalesced key
  // sits in its left position but draws each cell from whichever side is
  // present (`coalesce_columns`); the matching right key column is then dropped
  // below.
  //
  // A join is a row rebuild, so each output column canonicalises onto the
  // backend its own gathered cells imply (`docs/performance.md` classifies the
  // paths): all-valid numeric cells land on `Numeric`, and a column that picks
  // up a null from an unmatched row's `None` is nullable and stays `Builtin`.
  // `gather_series_opt` already does that — the `preserve_backend` wrapper is
  // the belt to its braces, a no-op on a canonicalised result either way, kept
  // so a future gather that stops canonicalising cannot silently demote a
  // `Numeric` source here. The coalesced key follows the **left** key column's
  // backend (its values share the pre-checked-equal key dtype).
  let out_cols : Array[Series] = []
  for col in self.column_series() {
    if coalesce && on_set.contains(col.name()) {
      // `other` has this key (validated above), so `get_column`'s
      // `ColumnNotFound` is forwarded but never taken.
      let left_key = gather_series_opt(col, left_rows, col.name())
      let right_key = gather_series_opt(
        other.get_column(col.name()),
        right_rows,
        col.name(),
      )
      // Surface the documented side's value: the left on `Inner` / `Left`, and
      // the present side on `Outer` (the left key is the base — taken where the
      // row has a left match, else the right fills the gap), but the **right**
      // on `Right`, where every output row has the right key present and the
      // left only on matched rows. Folded float keys make the side observable:
      // `-0.0` and `+0.0` share a group, so a `Right` join must surface the
      // right's signed zero, not the left's.
      let merged = match options.how {
        Right => coalesce_columns(right_key, left_key, col.name())
        Inner | Left | Outer | Cross =>
          coalesce_columns(left_key, right_key, col.name())
      }
      out_cols.push(preserve_backend(col, merged))
    } else {
      out_cols.push(
        preserve_backend(col, gather_series_opt(col, left_rows, col.name())),
      )
    }
  }
  // Right columns are gathered with `None -> null` and renamed on a name
  // collision with the left frame. A right key column is dropped only when
  // coalescing (its values are already in the kept key); otherwise every
  // right column is emitted, suffixed on a clash — a kept key column always
  // clashes, so it becomes ``.
  for col in other.column_series() {
    if !(coalesce && on_set.contains(col.name())) {
      let out_name = if left_names.contains(col.name()) {
        col.name() + options.suffix
      } else {
        col.name()
      }
      out_cols.push(
        preserve_backend(col, gather_series_opt(col, right_rows, out_name)),
      )
    }
  }
  // The row plan fixes the height: every gathered column is `left_rows` long
  // (`right_rows` is parallel to it), and when neither frame has a column to
  // infer from, the plan is the only thing that still knows how many rows the
  // join produced — `2×0` cross `3×0` is `6×0`.
  DataFrame::from_parts(out_cols, left_rows.length())
}

///|
/// Resolve and materialise the join key columns of both frames. Picks the
/// effective key expressions (`join_key_exprs`: `on` applied to both sides,
/// or the paired `left_on` / `right_on`), evaluates each over its own frame
/// (broadcasting a length-1 key up to frame height, like `sort` /
/// `group_by`), checks the left and right dtypes agree per key (so the
/// values could compare equal), and returns the per-key left and right cell
/// vectors (`to_scalars`, total — callers index by plain `[i]`). Shared by
/// the `Right` and `Inner` / `Left` / `Outer` row planners.
///
/// Raises:
/// - `InvalidOperation(detail)` — no keys at all, both `on` and
///   `left_on` / `right_on` given, or `left_on` / `right_on` of unequal
///   length (all via `join_key_exprs`).
/// - `DuplicateColumn(name)` — a bare `col(name)` key is repeated on an `on`
///   join. Reported at the repeat, in key order, mirroring `group_by`'s
///   key-list contract (a repeated bare `on` key would otherwise be silently
///   redundant). A `left_on` / `right_on` join does *not* check this — a
///   repeated left key is a valid multi-condition join (`a == b AND a == c`).
///   Only bare-`col` keys are checked: derived keys contribute no output
///   column, so two distinct derived keys sharing a leftmost name are different
///   keys, not a collision. Checked before that key is evaluated, so a
///   *missing* repeated key still surfaces as `ColumnNotFound` at its first
///   appearance.
/// - `ColumnNotFound(name)` — a key expression references an absent column,
///   first offending key in key order (the left key evaluated before the
///   right).
/// - `TypeMismatch(detail)` — a key whose left and right dtypes differ (or
///   whose own derived dtypes don't unify), first such key in key order.
/// - `LengthMismatch` — a `lit_series(s)` key whose embedded series is
///   neither length 1 nor the frame's height (the evaluator's broadcast
///   rule).
fn resolve_join_keys(
  left : DataFrame,
  right : DataFrame,
  options : JoinOptions,
) -> (Array[Array[@types.Scalar]], Array[Array[@types.Scalar]]) raise @types.DataError {
  let (left_exprs, right_exprs) = join_key_exprs(options)
  if left_exprs.is_empty() {
    raise @types.DataError::InvalidOperation(
      "a non-cross join requires at least one key column; use how = Cross for a Cartesian product",
    )
  }
  let ln = left.nrows()
  let rn = right.nrows()
  let lscope = Array::makei(ln, i => i)
  let rscope = Array::makei(rn, i => i)
  let left_key_scalars : Array[Array[@types.Scalar]] = []
  let right_key_scalars : Array[Array[@types.Scalar]] = []
  // The running `seen` set raises `DuplicateColumn` at a bare-`col` key
  // repeated by name — before that key is evaluated a second time — so an `on`
  // join matches `group_by`'s "no duplicate keys" contract instead of silently
  // treating two equal keys as one. This applies to `on` only (`dedup`): with
  // `left_on` / `right_on`, a repeated *left* key is a legitimate
  // multi-condition join (`left_on([a, a]).right_on([b, c])` means `a == b AND
  // a == c`, which Polars allows) — the left and right key columns carry their
  // own names, so there is no output-column collision to pre-empt. Only a bare
  // `col(name)` participates even on the `on` path: a *derived* key (any
  // non-`Col` expression) contributes no output column — the output is the two
  // frames' own columns, never the key expressions — so two distinct derived
  // keys that share a leftmost-column name (`col("ts") / day` and
  // `col("ts") % day`) are legitimately different keys and must not collide.
  // Checked before evaluation, so a *missing* repeated key still surfaces as
  // `ColumnNotFound` at its first appearance (its own evaluation fails before
  // the repeat is reached).
  let dedup = options.left_on.is_empty()
  let seen : Map[String, Unit] = Map([])
  for i in 0.. (Array[@expr.Expr], Array[@expr.Expr]) raise @types.DataError {
  let has_on = !options.on.is_empty()
  let has_sided = !options.left_on.is_empty() || !options.right_on.is_empty()
  if has_on && has_sided {
    raise @types.DataError::InvalidOperation(
      "join takes either `on` or `left_on` / `right_on`, not both",
    )
  }
  if has_on {
    (options.on, options.on)
  } else {
    if options.left_on.length() != options.right_on.length() {
      raise @types.DataError::InvalidOperation(
        "left_on and right_on must name the same number of keys; got \{options.left_on.length()} and \{options.right_on.length()}",
      )
    }
    (options.left_on, options.right_on)
  }
}

///|
/// The bare-column key names eligible for coalescing, or `None` when this
/// join cannot coalesce. Only an `on` join whose every key is a bare
/// `col(...)` reads the same-named column from both frames — the
/// precondition for merging a key into one output column. `left_on` /
/// `right_on` (paired, possibly differently-named keys) and any derived
/// `on` key turn coalescing off (matching Polars: "join on a non-column
/// expression turns off coalescing"), so both key columns are kept.
fn coalesce_key_names(options : JoinOptions) -> Array[String]? {
  if !options.left_on.is_empty() || !options.right_on.is_empty() {
    return None
  }
  let names = []
  for k in options.on {
    match k.node() {
      @ir.ExprNode::Col(name) => names.push(name)
      _ => return None
    }
  }
  Some(names)
}