aedb 0.3.1

//! Window-bounded join execution: the streaming PK-probe fast path that stops
//! reading once the page window is full, and the shared pagination /
//! projection / cursor tail used by both the streaming and general join paths.

use super::join::{
    MAX_JOIN_MATERIALIZED_BYTES, enforce_materialization_budget, is_single_column_primary_key_join,
    order_join_rows, project_join_rows, resolve_table_ref, simple_equi_columns,
};
use super::pagination::{
    compute_page_window, compute_remaining_limit_after_page, compute_split_recommended,
};
use super::pk_ordered_scan::{cursor_value_matches_pk_type, pk_encoded_order_matches_value_order};
use super::{
    CursorToken, QueryResult, encode_cursor, extract_pk_key, extract_sort_key, row_after_cursor,
};
use crate::catalog::Catalog;
use crate::catalog::namespace_key;
use crate::catalog::types::{ColumnType, Row, Value};
use crate::query::error::QueryError;
use crate::query::plan::{JoinType, Order, Query, QueryOptions};
use crate::storage::encoded_key::EncodedKey;
use crate::storage::keyspace::KeyspaceSnapshot;
use std::ops::Bound;

/// Shared tail of join execution: order, deduplicate (DISTINCT), paginate,
/// project, and encode the continuation cursor over the joined row set.
pub(super) struct FinalizeJoinRowsRequest<'a> {
    pub(super) rows: Vec<Row>,
    pub(super) columns: &'a [String],
    pub(super) query: &'a Query,
    pub(super) cursor_state: Option<CursorToken>,
    pub(super) snapshot_seq: u64,
    pub(super) max_scan_rows: usize,
    pub(super) cursor_signing_key: Option<&'a [u8; 32]>,
    pub(super) rows_examined: usize,
}

pub(super) fn finalize_join_rows(
    request: FinalizeJoinRowsRequest<'_>,
) -> Result<QueryResult, QueryError> {
    let FinalizeJoinRowsRequest {
        mut rows,
        columns,
        query,
        cursor_state,
        snapshot_seq,
        max_scan_rows,
        cursor_signing_key,
        rows_examined,
    } = request;
    let page_window = compute_page_window(query, &cursor_state, max_scan_rows)?;
    // When the result is ordered and capped (no cursor, no DISTINCT — both need
    // every row), only the first `offset + page + lookahead` rows are ever
    // consumed, so a bounded partial sort (top-k) avoids fully sorting the set.
    let topk_bound = if cursor_state.is_none()
        && !query.distinct
        && query.limit.is_some()
        && !query.order_by.is_empty()
    {
        Some(
            page_window
                .row_offset_count
                .saturating_add(page_window.page_read_limit),
        )
    } else {
        None
    };
    order_join_rows(&mut rows, columns, query, topk_bound)?;

    if query.distinct {
        // Project, then deduplicate the output rows before applying offset/limit.
        // Cursor pagination with DISTINCT is rejected before execution.
        let projected = project_join_rows(rows, columns, query)?;
        let mut seen: std::collections::HashSet<Vec<Value>> = std::collections::HashSet::new();
        let mut distinct_rows: Vec<Row> = Vec::new();
        for row in projected {
            if seen.insert(row.values.clone()) {
                distinct_rows.push(row);
            }
        }
        let mut sliced: Vec<Row> = distinct_rows
            .into_iter()
            .skip(page_window.row_offset_count)
            .collect();
        let has_more = sliced.len() > page_window.effective_page_size;
        if has_more {
            sliced.truncate(page_window.effective_page_size);
        }
        let split_budget = query.limit.unwrap_or(max_scan_rows);
        let split_recommended = compute_split_recommended(rows_examined, split_budget);
        return Ok(QueryResult {
            rows_examined,
            rows: sliced,
            truncated: has_more,
            cursor: None,
            snapshot_seq,
            materialized_seq: None,
            split_recommended,
        });
    }

    let sort_indices: Vec<(usize, crate::query::plan::Order)> = if !query.order_by.is_empty() {
        query
            .order_by
            .iter()
            .filter_map(|(name, ord)| columns.iter().position(|c| c == name).map(|i| (i, *ord)))
            .collect()
    } else {
        Vec::new()
    };
    let pk_indices: Vec<usize> = (0..columns.len()).collect();
    let mut sliced = Vec::new();
    let mut skipped = 0usize;
    for row in rows {
        if let Some(cursor) = &cursor_state
            && !row_after_cursor(&row, cursor, &sort_indices, &pk_indices)
        {
            continue;
        }
        if skipped < page_window.row_offset_count {
            skipped += 1;
            continue;
        }
        sliced.push(row);
        if sliced.len() > page_window.effective_page_size {
            break;
        }
    }
    let has_more = sliced.len() > page_window.effective_page_size;
    if has_more {
        sliced.truncate(page_window.effective_page_size);
    }
    let cursor_last_row = sliced.last().cloned();

    sliced = project_join_rows(sliced, columns, query)?;

    let returned_rows = sliced.len();
    let remaining_limit_after_page = compute_remaining_limit_after_page(
        query.limit,
        cursor_state.as_ref().and_then(|c| c.remaining_limit),
        returned_rows,
    );
    let cursor = if has_more && page_window.row_offset_count == 0 {
        let last_row = cursor_last_row.ok_or_else(|| QueryError::InvalidQuery {
            reason: "invalid cursor state".into(),
        })?;
        Some(encode_cursor(
            &CursorToken {
                snapshot_seq,
                last_sort_key: extract_sort_key(&last_row, &sort_indices),
                last_pk: extract_pk_key(&last_row, &pk_indices),
                page_size: page_window.page_size,
                remaining_limit: remaining_limit_after_page,
            },
            cursor_signing_key,
        )?)
    } else {
        None
    };
    let split_budget = query.limit.unwrap_or(max_scan_rows);
    let split_recommended = compute_split_recommended(rows_examined, split_budget);
    Ok(QueryResult {
        rows_examined,
        rows: sliced,
        truncated: has_more,
        cursor,
        snapshot_seq,
        materialized_seq: None,
        split_recommended,
    })
}

/// Joined rows produced by the streaming PK-probe fast path: already in output
/// order, bounded to the page window, with the combined column layout.
pub(super) struct StreamedPkProbeJoin {
    pub(super) rows: Vec<Row>,
    pub(super) columns: Vec<String>,
    pub(super) rows_examined: usize,
}

pub(super) struct StreamingPkProbeJoinRequest<'a> {
    pub(super) snapshot: &'a KeyspaceSnapshot,
    pub(super) catalog: &'a Catalog,
    pub(super) project_id: &'a str,
    pub(super) scope_id: &'a str,
    pub(super) query: &'a Query,
    pub(super) options: &'a QueryOptions,
    pub(super) cursor_state: &'a Option<CursorToken>,
    pub(super) max_scan_rows: usize,
    pub(super) base_schema: &'a crate::catalog::schema::TableSchema,
    pub(super) base_ns_project: &'a str,
    pub(super) base_ns_scope: &'a str,
    pub(super) base_table: &'a str,
    pub(super) base_alias: &'a str,
    pub(super) base_columns: &'a [String],
    pub(super) base_column_types: &'a [ColumnType],
    pub(super) base_count: usize,
}

/// Stream a single PK-probe equi join in base primary-key order, stopping once
/// the page window (`offset + limit + lookahead`) is full — instead of
/// materializing and joining both tables and discarding everything past the
/// page. Returns `Ok(None)` when the query shape does not guarantee that the
/// streamed prefix equals what full evaluation would paginate to:
///
/// - exactly one INNER/LEFT join whose right key is the join table's
///   single-column primary key with exactly matching column types (the same
///   gate as the non-streaming PK-probe path);
/// - no WHERE / aggregates / GROUP BY / HAVING / DISTINCT (each needs the full
///   joined set or interacts with pushdown), and an explicit LIMIT (without
///   one, `rows_examined` must reflect full evaluation);
/// - ORDER BY absent (output order is base-scan order either way) or exactly
///   `base_pk ASC|DESC` on an order-preserving PK type — DESC streams the
///   base map in reverse; a cursor additionally requires the ORDER BY form so
///   `row_after_cursor` is monotone in stream order;
/// - the streamed base table fully resident (no cold row segments or
///   tombstones), so its hot PK map is the complete iteration source. Join
///   tables carry no residency gate: probes are tier-aware per key.
pub(super) fn try_streaming_pk_probe_join(
    req: StreamingPkProbeJoinRequest<'_>,
) -> Result<Option<StreamedPkProbeJoin>, QueryError> {
    if req.query.joins.len() != 1
        || req.query.predicate.is_some()
        || !req.query.aggregates.is_empty()
        || !req.query.group_by.is_empty()
        || req.query.having.is_some()
        || req.query.distinct
        || req.query.limit.is_none()
    {
        return Ok(None);
    }
    let join = &req.query.joins[0];
    if !matches!(join.join_type, JoinType::Inner | JoinType::Left) {
        return Ok(None);
    }
    let mut base_pk_type: Option<&ColumnType> = None;
    let mut reverse = false;
    match req.query.order_by.as_slice() {
        [] => {
            if req.cursor_state.is_some() {
                return Ok(None);
            }
        }
        [(col, order)] => {
            if req.base_schema.primary_key.len() != 1 {
                return Ok(None);
            }
            let pk_name = &req.base_schema.primary_key[0];
            if *col != format!("{}.{pk_name}", req.base_alias) {
                return Ok(None);
            }
            let pk_type = req
                .base_schema
                .columns
                .iter()
                .find(|c| &c.name == pk_name)
                .map(|c| &c.col_type);
            if !pk_encoded_order_matches_value_order(pk_type) {
                return Ok(None);
            }
            base_pk_type = pk_type;
            reverse = matches!(order, Order::Desc);
        }
        _ => return Ok(None),
    }

    let (jp, js, jt) = resolve_table_ref(req.project_id, req.scope_id, &join.table);
    let Some(join_schema) = req
        .catalog
        .tables
        .get(&(namespace_key(&jp, &js), jt.clone()))
    else {
        // Let the general path surface TableNotFound.
        return Ok(None);
    };
    let join_alias = join.alias.clone().unwrap_or(jt.clone());
    let join_qualified_cols: Vec<String> = join_schema
        .columns
        .iter()
        .map(|c| format!("{join_alias}.{}", c.name))
        .collect();
    let simple_equi = join
        .on
        .as_ref()
        .and_then(|on| simple_equi_columns(on, req.base_columns, &join_qualified_cols));
    if join.on.is_some() && simple_equi.is_none() {
        return Ok(None);
    }
    let effective_left = match simple_equi
        .as_ref()
        .map(|(l, _)| l.clone())
        .or_else(|| join.left_column.clone())
    {
        Some(column) => column,
        None => return Ok(None),
    };
    let effective_right = match simple_equi
        .as_ref()
        .map(|(_, r)| r.clone())
        .or_else(|| join.right_column.clone())
    {
        Some(column) => column,
        None => return Ok(None),
    };
    // Unresolvable join columns fall back so the general path raises the same
    // ColumnNotFound errors it does today.
    let Some(left_idx) = req.base_columns.iter().position(|c| *c == effective_left) else {
        return Ok(None);
    };
    let Some(right_idx) = join_schema.columns.iter().position(|c| {
        format!("{join_alias}.{}", c.name) == effective_right || c.name == effective_right
    }) else {
        return Ok(None);
    };
    if req.base_column_types.get(left_idx)
        != join_schema.columns.get(right_idx).map(|c| &c.col_type)
    {
        return Ok(None);
    }
    if !is_single_column_primary_key_join(join_schema, right_idx) {
        return Ok(None);
    }

    // The streamed BASE table must be fully resident: with cold row segments
    // or tombstones its hot PK map is an incomplete iteration source, so
    // defer to the general path. The JOIN table needs no such gate — probes
    // go through the tier-aware row lookup, which resolves cold copies and
    // tombstones per key.
    let base_table_data =
        req.snapshot
            .table(req.base_ns_project, req.base_ns_scope, req.base_table);
    if base_table_data
        .map(|t| !t.row_segments.is_empty() || !t.row_tombstones.is_empty())
        .unwrap_or(false)
    {
        return Ok(None);
    }
    let join_table_data = req.snapshot.table(&jp, &js, &jt);

    // Same base-table scan guard as the general full-scan path (this shape has
    // no pushdown predicate, so the general path would always full-scan).
    if !req.options.allow_full_scan && req.base_count > req.max_scan_rows {
        return Err(QueryError::ScanBoundExceeded {
            estimated_rows: req.base_count as u64,
            max_scan_rows: req.max_scan_rows as u64,
        });
    }

    let page_window = compute_page_window(req.query, req.cursor_state, req.max_scan_rows)?;
    let window_limit = page_window.row_source_window_limit;
    let mut columns: Vec<String> = req.base_columns.to_vec();
    columns.extend(join_qualified_cols);
    let sort_indices: Vec<(usize, Order)> = req
        .query
        .order_by
        .iter()
        .filter_map(|(name, ord)| columns.iter().position(|c| c == name).map(|i| (i, *ord)))
        .collect();
    let pk_indices: Vec<usize> = (0..columns.len()).collect();
    let join_col_count = join_schema.columns.len();
    let preserve_left = matches!(join.join_type, JoinType::Left);

    // Resume directly past the cursor position instead of skip-scanning from
    // the start: the stream order is the base PK, so the cursor's sort key
    // (the last returned base PK value) locates the restart point — as an
    // exclusive lower bound ascending, or an exclusive upper bound descending.
    // Only taken when that value's type exactly matches the PK column — then
    // encoded-key order equals `row_after_cursor`'s `Value::cmp` order, so
    // the seek can never jump past an admissible row. The per-row cursor
    // filter below stays either way, as the semantic gate.
    let bounds: (Bound<EncodedKey>, Bound<EncodedKey>) = match (req.cursor_state, base_pk_type) {
        (Some(cursor), Some(pk_type))
            if cursor.last_sort_key.len() == 1
                && cursor_value_matches_pk_type(&cursor.last_sort_key[0], pk_type) =>
        {
            let key = EncodedKey::from_values(&cursor.last_sort_key);
            if reverse {
                (Bound::Unbounded, Bound::Excluded(key))
            } else {
                (Bound::Excluded(key), Bound::Unbounded)
            }
        }
        _ => (Bound::Unbounded, Bound::Unbounded),
    };

    let mut rows: Vec<Row> = Vec::new();
    let mut rows_examined = 0usize;
    if let Some(base) = base_table_data {
        let range = base.rows.range(bounds);
        let stored_rows: Box<dyn Iterator<Item = &crate::storage::keyspace::StoredRow>> = if reverse
        {
            Box::new(range.rev().map(|(_, stored)| stored))
        } else {
            Box::new(range.map(|(_, stored)| stored))
        };
        for stored in stored_rows {
            if rows.len() >= window_limit {
                break;
            }
            let left = req.snapshot.materialize_row(stored)?;
            let probe = EncodedKey::from_values(std::slice::from_ref(&left.values[left_idx]));
            // Tier-aware probe: cold-tier join-table copies resolve per key.
            let matched = match join_table_data {
                Some(table) => req.snapshot.get_row_in_table(table, &probe)?,
                None => None,
            };
            let joined = match matched {
                Some(right) => {
                    let mut values =
                        Vec::with_capacity(left.values.len().saturating_add(right.values.len()));
                    values.extend_from_slice(&left.values);
                    values.extend_from_slice(&right.values);
                    Row { values }
                }
                None if preserve_left => {
                    let mut values =
                        Vec::with_capacity(left.values.len().saturating_add(join_col_count));
                    values.extend_from_slice(&left.values);
                    values.extend(std::iter::repeat_n(Value::Null, join_col_count));
                    Row { values }
                }
                None => continue,
            };
            rows_examined += 1;
            if let Some(cursor) = req.cursor_state
                && !row_after_cursor(&joined, cursor, &sort_indices, &pk_indices)
            {
                continue;
            }
            rows.push(joined);
        }
    }
    enforce_materialization_budget(&rows, MAX_JOIN_MATERIALIZED_BYTES)?;
    Ok(Some(StreamedPkProbeJoin {
        rows,
        columns,
        rows_examined,
    }))
}