vibesql_executor/evaluator/window/
partitioning.rs

1//! Row partitioning for window functions
2//!
3//! Groups rows into partitions based on PARTITION BY expressions.
4
5use vibesql_ast::Expression;
6use vibesql_storage::Row;
7use vibesql_types::SqlValue;
8
9/// A partition of rows for window function evaluation
10#[derive(Debug, Clone)]
11pub struct Partition {
12    pub rows: Vec<Row>,
13    /// Original indices of rows before partitioning/sorting
14    pub original_indices: Vec<usize>,
15    /// Column name to index mapping (for resolving named column references)
16    pub column_map: std::collections::HashMap<String, usize>,
17}
18
19impl Partition {
20    pub fn new(rows: Vec<Row>) -> Self {
21        let original_indices = (0..rows.len()).collect();
22        Self { rows, original_indices, column_map: std::collections::HashMap::new() }
23    }
24
25    pub fn with_indices(rows: Vec<Row>, original_indices: Vec<usize>) -> Self {
26        Self { rows, original_indices, column_map: std::collections::HashMap::new() }
27    }
28
29    pub fn with_column_map(
30        rows: Vec<Row>,
31        original_indices: Vec<usize>,
32        column_map: std::collections::HashMap<String, usize>,
33    ) -> Self {
34        Self { rows, original_indices, column_map }
35    }
36
37    pub fn len(&self) -> usize {
38        self.rows.len()
39    }
40
41    pub fn is_empty(&self) -> bool {
42        self.rows.is_empty()
43    }
44
45    /// Get column index by name (case-insensitive)
46    pub fn get_column_index(&self, name: &str) -> Option<usize> {
47        // First try exact match
48        if let Some(&idx) = self.column_map.get(name) {
49            return Some(idx);
50        }
51        // Try lowercase match
52        let lower = name.to_lowercase();
53        self.column_map.get(&lower).copied()
54    }
55}
56
57/// Partition rows by PARTITION BY expressions
58///
59/// Groups rows into partitions based on partition expressions.
60/// If no PARTITION BY clause, all rows go into a single partition.
61///
62/// Partitions are ordered by their key values (using BTreeMap with SqlValue keys
63/// for proper numeric/string ordering), matching SQLite's behavior.
64///
65/// Returns an error if evaluating any PARTITION BY expression fails. Expressions
66/// that legitimately evaluate to NULL still group together under the NULL key.
67pub fn partition_rows<F>(
68    rows: Vec<Row>,
69    partition_by: &Option<Vec<Expression>>,
70    eval_fn: F,
71) -> Result<Vec<Partition>, String>
72where
73    F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
74{
75    partition_rows_with_collations(rows, partition_by, &[], eval_fn)
76}
77
78/// Partition rows by PARTITION BY expressions, respecting column collations
79///
80/// `collations` maps each PARTITION BY expression index to its resolved collation.
81/// For NOCASE collation, partition keys are compared case-insensitively.
82///
83/// Returns an error if evaluating any PARTITION BY expression fails. Expressions
84/// that legitimately evaluate to NULL still group together under the NULL key.
85pub fn partition_rows_with_collations<F>(
86    rows: Vec<Row>,
87    partition_by: &Option<Vec<Expression>>,
88    collations: &[Option<String>],
89    eval_fn: F,
90) -> Result<Vec<Partition>, String>
91where
92    F: Fn(&Expression, &Row) -> Result<SqlValue, String>,
93{
94    // If no PARTITION BY, return all rows in single partition
95    let Some(partition_exprs) = partition_by else {
96        return Ok(vec![Partition::new(rows)]);
97    };
98
99    if partition_exprs.is_empty() {
100        return Ok(vec![Partition::new(rows)]);
101    }
102
103    // Group rows by partition key values
104    // Use BTreeMap with SqlValue keys for proper value ordering (not string ordering)
105    // This ensures Integer(3) < Integer(7) < Integer(11) as expected
106    let mut partitions_map: std::collections::BTreeMap<Vec<SqlValue>, Vec<(usize, Row)>> =
107        std::collections::BTreeMap::new();
108
109    for (original_idx, row) in rows.into_iter().enumerate() {
110        // Evaluate partition expressions for this row
111        let mut partition_key = Vec::new();
112
113        for (i, expr) in partition_exprs.iter().enumerate() {
114            // Propagate evaluation errors instead of collapsing them into a NULL
115            // partition key, which would silently merge all affected rows into
116            // one partition and produce wrong window results.
117            let value = eval_fn(expr, &row)?;
118            // Apply collation to partition key for case-insensitive grouping
119            let collation = collations.get(i).and_then(|c| c.as_deref());
120            if let Some(coll) = collation {
121                if coll.eq_ignore_ascii_case("nocase") {
122                    let normalized = match &value {
123                        SqlValue::Varchar(s) => {
124                            SqlValue::Varchar(arcstr::ArcStr::from(s.to_uppercase()))
125                        }
126                        SqlValue::Character(s) => {
127                            SqlValue::Character(arcstr::ArcStr::from(s.to_uppercase()))
128                        }
129                        other => other.clone(),
130                    };
131                    partition_key.push(normalized);
132                    continue;
133                }
134            }
135            partition_key.push(value);
136        }
137
138        partitions_map.entry(partition_key).or_default().push((original_idx, row));
139    }
140
141    // Convert BTreeMap to Vec<Partition>, preserving sorted partition key order
142    Ok(partitions_map
143        .into_values()
144        .map(|rows_with_indices| {
145            let (indices, rows): (Vec<_>, Vec<_>) = rows_with_indices.into_iter().unzip();
146            Partition::with_indices(rows, indices)
147        })
148        .collect())
149}
vibesql_executor/evaluator/window/partitioning.rs

vibesql_executor/evaluator/window/
partitioning.rs
