Skip to main content

cqlite_core/schema/
key_ordering.rs

1//! Ordered partition/clustering key accessors for [`TableSchema`] (issue #1677).
2//!
3//! Extracted from `schema/mod.rs` (source-split doctrine, issue #1116) so the
4//! per-call re-sort that finding R6 flagged can be removed without growing the
5//! already-over-threshold `mod.rs`.
6//!
7//! [`TableSchema::ordered_partition_keys`] / [`TableSchema::ordered_clustering_keys`]
8//! previously sorted their key slice by `position` on **every** call. During
9//! CQL→mutation building these are hit once per statement, so a `BEGIN BATCH` of N
10//! statements re-sorted N times even though the schema — and hence the ordering —
11//! is fixed. Key `position`s are assigned contiguously in order at schema
12//! construction (see `TableSchema::from_sstable_header`), so the slice is virtually
13//! always already sorted; the fast path below detects that with one linear scan and
14//! skips the sort entirely. When a schema *is* built out of order the sort still
15//! runs, so the returned order is byte-for-byte identical to before.
16
17use super::{ClusteringColumn, KeyColumn, TableSchema};
18
19impl TableSchema {
20    /// Get partition key columns ordered by `position`.
21    ///
22    /// Returns the same order as a full `sort_by_key(position)` would, but skips
23    /// the sort when the keys are already in ascending `position` order (the
24    /// common case — see the module docs), so it is no longer re-sorted per call.
25    pub fn ordered_partition_keys(&self) -> Vec<&KeyColumn> {
26        let mut keys = self.partition_keys.iter().collect::<Vec<_>>();
27        if !keys.is_sorted_by_key(|k| k.position) {
28            keys.sort_by_key(|k| k.position);
29        }
30        keys
31    }
32
33    /// Get clustering key columns ordered by `position`.
34    ///
35    /// Same order as before; the sort is skipped when the keys are already ordered.
36    pub fn ordered_clustering_keys(&self) -> Vec<&ClusteringColumn> {
37        let mut keys = self.clustering_keys.iter().collect::<Vec<_>>();
38        if !keys.is_sorted_by_key(|k| k.position) {
39            keys.sort_by_key(|k| k.position);
40        }
41        keys
42    }
43}
44
45#[cfg(test)]
46mod tests {
47    use super::*;
48    use crate::schema::{ClusteringColumn, ClusteringOrder, KeyColumn};
49    use std::collections::HashMap;
50
51    fn schema_with(
52        partition_keys: Vec<KeyColumn>,
53        clustering_keys: Vec<ClusteringColumn>,
54    ) -> TableSchema {
55        TableSchema {
56            keyspace: "ks".into(),
57            table: "t".into(),
58            partition_keys,
59            clustering_keys,
60            columns: vec![],
61            comments: HashMap::new(),
62            dropped_columns: HashMap::new(),
63        }
64    }
65
66    fn pk(name: &str, position: usize) -> KeyColumn {
67        KeyColumn {
68            name: name.into(),
69            data_type: "int".into(),
70            position,
71        }
72    }
73
74    fn ck(name: &str, position: usize) -> ClusteringColumn {
75        ClusteringColumn {
76            name: name.into(),
77            data_type: "int".into(),
78            position,
79            order: ClusteringOrder::Asc,
80        }
81    }
82
83    #[test]
84    fn already_sorted_keys_keep_their_order() {
85        let schema = schema_with(
86            vec![pk("a", 0), pk("b", 1), pk("c", 2)],
87            vec![ck("x", 0), ck("y", 1)],
88        );
89        let pks: Vec<_> = schema
90            .ordered_partition_keys()
91            .iter()
92            .map(|k| k.name.clone())
93            .collect();
94        assert_eq!(pks, vec!["a", "b", "c"]);
95        let cks: Vec<_> = schema
96            .ordered_clustering_keys()
97            .iter()
98            .map(|k| k.name.clone())
99            .collect();
100        assert_eq!(cks, vec!["x", "y"]);
101    }
102
103    #[test]
104    fn out_of_order_keys_are_still_sorted_by_position() {
105        // Byte-identical to the previous unconditional sort: order is by position,
106        // not by insertion order.
107        let schema = schema_with(
108            vec![pk("c", 2), pk("a", 0), pk("b", 1)],
109            vec![ck("y", 1), ck("x", 0)],
110        );
111        let pks: Vec<_> = schema
112            .ordered_partition_keys()
113            .iter()
114            .map(|k| (k.name.clone(), k.position))
115            .collect();
116        assert_eq!(pks, vec![("a".into(), 0), ("b".into(), 1), ("c".into(), 2)]);
117        let cks: Vec<_> = schema
118            .ordered_clustering_keys()
119            .iter()
120            .map(|k| (k.name.clone(), k.position))
121            .collect();
122        assert_eq!(cks, vec![("x".into(), 0), ("y".into(), 1)]);
123    }
124
125    #[test]
126    fn empty_key_lists_are_handled() {
127        let schema = schema_with(vec![], vec![]);
128        assert!(schema.ordered_partition_keys().is_empty());
129        assert!(schema.ordered_clustering_keys().is_empty());
130    }
131}