perf: cover scans with primary keys and estimate key prefixes for TPC-C

Cargo.toml

@@ -54,6 +54,7 @@ required-features = ["pprof"]
 
 [dependencies]
 bumpalo               = { version = "3", default-features = false, features = ["collections"] }
+hashbrown             = { version = "0.15", default-features = false }
 ordered-float         = { version = "4" }
 paste                 = { version = "1" }
 kite_sql_serde_macros = { version = "0.2.6", path = "kite_sql_serde_macros" }

README.md

@@ -135,9 +135,9 @@ print(list(db.run("select * from demo")))
 
 | Backend | TpmC | New-Order | Payment | Order-Status | Delivery | Stock-Level |
 | --- | ---: | ---: | ---: | ---: | ---: | ---: |
-| KiteSQL LMDB | 187161 | 274 | 68 | 40 | 454 | 278 |
-| KiteSQL RocksDB | 46291 | 374 | 184 | 201 | 14463 | 651 |
-| SQLite (WAL, `synchronous=OFF`) | 179039 | 291 | 43 | 38 | 368 | 320 |
+| KiteSQL LMDB | 209840 | 254 | 50 | 28 | 345 | 202 |
+| KiteSQL RocksDB | 48049 | 354 | 85 | 146 | 14511 | 562 |
+| SQLite (WAL, `synchronous=OFF`) | 192435 | 270 | 40 | 33 | 344 | 265 |
 
 👉 [Details and how to reproduce](tpcc/README.md)
 

src/db/program.rs

@@ -218,7 +218,7 @@ mod tests {
     use super::*;
     use crate::catalog::{ColumnCatalog, ColumnDesc};
     use crate::execution::dql::index_scan::IndexScan;
-    use crate::expression::range_detacher::Range;
+    use crate::expression::range_detacher::{Range, RangeValue};
     use crate::expression::{BinaryOperator, ScalarExpression};
     use crate::planner::operator::filter::FilterOperator;
     use crate::planner::operator::sort::SortField;
@@ -287,11 +287,11 @@ mod tests {
         });
         let original = Range::Scope {
             min: Bound::Unbounded,
-            max: Bound::Excluded(DataValue::Int32(5)),
+            max: Bound::Excluded(DataValue::Int32(5).into()),
         };
         let expected = Range::Scope {
-            min: Bound::Included(DataValue::Int32(3)),
-            max: Bound::Excluded(DataValue::Int32(5)),
+            min: Bound::Included(DataValue::Int32(3).into()),
+            max: Bound::Excluded(DataValue::Int32(5).into()),
         };
         let sort = SortOption::OrderBy {
             fields: vec![SortField::new(col_expr, true, false)],
@@ -349,7 +349,7 @@ mod tests {
                 lookup: Some(lookup),
                 residual_predicate: has_residual.then_some(predicate),
                 sort_option: sort.clone(),
-                covered_deserializers: None,
+                covered: false,
                 cover_mapping: None,
                 sort_elimination_hint: None,
                 stream_aggregate_hint: None,
@@ -600,21 +600,61 @@ mod tests {
         Ok(())
     }
 
-    #[test]
-    fn composite_index_ranges_bind_for_each_execution() -> Result<(), DatabaseError> {
-        fn index_range(plan: &LogicalPlan) -> Option<&Range> {
-            plan.physical_option
-                .as_ref()
-                .and_then(|option| match &option.plan {
-                    PlanImpl::IndexScan(info) => match &info.lookup {
-                        Some(IndexLookup::Static(range)) => Some(range),
-                        _ => None,
-                    },
+    fn index_range(plan: &LogicalPlan) -> Option<&Range> {
+        plan.physical_option
+            .as_ref()
+            .and_then(|option| match &option.plan {
+                PlanImpl::IndexScan(info) => match &info.lookup {
+                    Some(IndexLookup::Static(range)) => Some(range),
                     _ => None,
-                })
-                .or_else(|| plan.childrens.iter().find_map(index_range))
+                },
+                _ => None,
+            })
+            .or_else(|| plan.childrens.iter().find_map(index_range))
+    }
+
+    #[test]
+    fn null_parameters_empty_index_ranges() -> Result<(), DatabaseError> {
+        let mut db = DataBaseBuilder::path(".")
+            .histogram_buckets(2)
+            .build_in_memory()?;
+        db.ddl("create table t(w int, k int, primary key(w,k))")?;
+        db.run("insert into t values(1,1),(1,2),(2,1),(2,2),(2,4)")?
+            .done()?;
+        db.analyze("t")?;
+        let ty = LogicalType::Integer;
+        let mut tx = db.new_transaction(false)?;
+        // A NULL bound used to reach storage as a NULL key: the composite scope below became
+        // `(2, NULL) .. (2, 5)`, an inverted range, and `w < NULL` scanned every row.
+        for (sql, params) in [
+            (
+                "select k from t where w = $1 and k < $2 and k >= $3",
+                vec![
+                    (1, DataValue::Int32(2)),
+                    (2, DataValue::Int32(5)),
+                    (3, DataValue::Null),
+                ],
+            ),
+            ("select k from t where w < $1", vec![(1, DataValue::Null)]),
+            ("select k from t where w = $1", vec![(1, DataValue::Null)]),
+        ] {
+            let mut plan = db.program_sql(
+                sql,
+                &(1..=params.len())
+                    .map(|id| (id, ty.clone()))
+                    .collect::<Vec<_>>(),
+            )?;
+            assert!(index_range(plan.executor.plan()).is_some(), "{sql}");
+            let mut iter = plan.execute(&mut tx, &params)?;
+            assert!(iter.next_tuple(|_, _| ())?.is_none(), "{sql}");
+            iter.done()?;
         }
 
+        Ok(())
+    }
+
+    #[test]
+    fn composite_index_ranges_bind_for_each_execution() -> Result<(), DatabaseError> {
         let mut db = DataBaseBuilder::path(".")
             .histogram_buckets(2)
             .build_in_memory()?;
@@ -694,6 +734,19 @@ mod tests {
         else {
             panic!("expected index scan")
         };
+        // `$3 - 20` is computed when the range is bound, so the whole predicate is the range.
+        fn has_filter(plan: &LogicalPlan) -> bool {
+            matches!(plan.operator, Operator::Filter(_)) || plan.childrens.iter().any(has_filter)
+        }
+        assert!(matches!(
+            &info.lookup,
+            Some(IndexLookup::Static(Range::Scope {
+                min: Bound::Included(RangeValue::Computed(_)),
+                ..
+            }))
+        ));
+        assert!(info.residual_predicate.is_none());
+        assert!(!has_filter(plan.executor.plan()));
         let tables = TableCache::default();
         let context = &tables;
         let storage = MemoryStorage::new();
@@ -712,16 +765,41 @@ mod tests {
         assert_eq!(
             ranges.next(),
             Some(&Range::Scope {
-                min: Bound::Included(DataValue::Tuple(vec![
-                    DataValue::Int32(2),
-                    DataValue::Int32(2),
-                ])),
-                max: Bound::Excluded(DataValue::Tuple(vec![
-                    DataValue::Int32(2),
-                    DataValue::Int32(5),
-                ])),
+                min: Bound::Included(
+                    DataValue::Tuple(vec![DataValue::Int32(2), DataValue::Int32(2),]).into()
+                ),
+                max: Bound::Excluded(
+                    DataValue::Tuple(vec![DataValue::Int32(2), DataValue::Int32(5),]).into()
+                ),
             })
         );
+        drop(ranges);
+
+        // The computed bound follows each execution's parameters, and a NULL one matches nothing.
+        let mut tx = db.new_transaction(false)?;
+        for (lower, expected) in [
+            (DataValue::Int32(22), vec![2, 4]),
+            (DataValue::Int32(24), vec![4]),
+            (DataValue::Null, vec![]),
+        ] {
+            let mut iter = plan.execute(
+                &mut tx,
+                [
+                    (1, DataValue::Int32(2)),
+                    (2, DataValue::Int32(5)),
+                    (3, lower),
+                ],
+            )?;
+            let mut rows = Vec::new();
+            while let Some(k) = iter.next_tuple(|_, row| row.value_at(0).clone())? {
+                rows.push(k);
+            }
+            let expected = expected
+                .into_iter()
+                .map(DataValue::Int32)
+                .collect::<Vec<_>>();
+            assert_eq!(rows, expected);
+        }
         Ok(())
     }
 }

src/execution/dml/analyze.rs

@@ -277,25 +277,28 @@ mod test {
     }
 
     fn test_hll_precision_is_configurable() -> Result<(), DatabaseError> {
-        let rows = 1000;
+        let (rows, distinct) = (1000, 500);
         let ndv_with = |precision: u8| -> Result<usize, DatabaseError> {
             let mut builder = HistogramMetaBuilder::new(precision)?;
             for i in 0..rows {
-                builder.append(&DataValue::Int32(i));
+                builder.append(&DataValue::Int32(i % distinct));
             }
             Ok(builder.build().distinct_values_len())
         };
         let (coarse, default) = (ndv_with(4)?, ndv_with(DEFAULT_HLL_PRECISION)?);
         // Precision 4 has 16 registers, so its estimate must differ from the default sketch.
         assert_ne!(coarse, default);
 
-        let values = (0..rows).map(|i| format!("({i})")).join(", ");
+        let values = (0..rows)
+            .map(|i| format!("({i}, {})", i % distinct))
+            .join(", ");
         for (builder, expected) in [
             (DataBaseBuilder::path("."), default),
             (DataBaseBuilder::path(".").hll_precision(4), coarse),
         ] {
             let mut kite_sql = builder.build_in_memory()?;
-            kite_sql.ddl("create table t1 (a int primary key)")?;
+            kite_sql.ddl("create table t1 (id int primary key, a int)")?;
+            kite_sql.ddl("create index idx_a on t1 (a)")?;
             kite_sql
                 .run(format!("insert into t1 values {values}"))?
                 .done()?;
@@ -318,7 +321,7 @@ mod test {
             let loader = StatisticMetaLoader::new(kite_sql.state.meta_cache());
             let column = loader.load_column(table_id, a_id).unwrap();
             assert_eq!(column.distinct_values_len(), expected);
-            let index = loader.load(table_id, 0)?.unwrap();
+            let index = loader.load(table_id, 1)?.unwrap();
             assert_eq!(index.histogram().distinct_values_len(), expected);
         }
         Ok(())
@@ -438,7 +441,7 @@ mod test {
         let expected_frequency = loader.collect_count(
             table_id,
             1,
-            &Range::Eq(DataValue::Int32(7)),
+            &Range::Eq(DataValue::Int32(7).into()),
             &mut Vec::new(),
         )?;
 
@@ -457,7 +460,7 @@ mod test {
             loader.collect_count(
                 table_id,
                 1,
-                &Range::Eq(DataValue::Int32(7)),
+                &Range::Eq(DataValue::Int32(7).into()),
                 &mut Vec::new()
             )?,
             expected_frequency
@@ -586,7 +589,7 @@ mod test {
                 .collect_count(
                     table_id,
                     1,
-                    &Range::Eq(DataValue::Int32(7)),
+                    &Range::Eq(DataValue::Int32(7).into()),
                     &mut Vec::new()
                 )?
                 .map(|(rows, _)| rows),
@@ -613,7 +616,7 @@ mod test {
                 .collect_count(
                     table_id,
                     1,
-                    &Range::Eq(DataValue::Int32(7)),
+                    &Range::Eq(DataValue::Int32(7).into()),
                     &mut Vec::new()
                 )?
                 .map(|(rows, _)| rows),

src/execution/dql/aggregate/avg.rs

@@ -15,15 +15,16 @@
 use crate::errors::DatabaseError;
 use crate::execution::dql::aggregate::sum::{text_as_number, SumAccumulator};
 use crate::execution::dql::aggregate::Accumulator;
+use crate::execution::hash::FxHashSet;
 use crate::expression::BinaryOperator;
 use crate::types::evaluator::binary_create;
 use crate::types::value::DataValue;
 use crate::types::LogicalType;
 use std::borrow::Cow;
-use std::collections::HashSet;
+use std::mem;
 
 pub struct AvgAccumulator {
-    distinct_values: Option<HashSet<DataValue>>,
+    distinct_values: Option<FxHashSet<DataValue>>,
     inner: Option<SumAccumulator>,
     count: usize,
     result: DataValue,
@@ -32,7 +33,7 @@ pub struct AvgAccumulator {
 impl AvgAccumulator {
     pub fn new(distinct: bool) -> Self {
         Self {
-            distinct_values: distinct.then(HashSet::new),
+            distinct_values: distinct.then(FxHashSet::default),
             inner: None,
             count: 0,
             result: DataValue::Null,
@@ -85,8 +86,17 @@ impl Accumulator for AvgAccumulator {
         &self.result
     }
 
-    fn result_owned(self: Box<Self>) -> DataValue {
-        self.result
+    fn take_result(&mut self) -> DataValue {
+        mem::replace(&mut self.result, DataValue::Null)
+    }
+
+    fn reset(&mut self) {
+        if let Some(values) = &mut self.distinct_values {
+            values.clear();
+        }
+        self.inner = None;
+        self.count = 0;
+        self.result = DataValue::Null;
     }
 }
 
@@ -109,7 +119,7 @@ mod tests {
         assert_ne!(first, second);
         accumulator.evaluate()?;
         assert_eq!(&second, accumulator.result());
-        assert_eq!(Box::new(accumulator).result_owned(), second);
+        assert_eq!(accumulator.take_result(), second);
         Ok(())
     }
 
@@ -133,10 +143,7 @@ mod tests {
 
         accumulator.update_value(&DataValue::Int32(3))?;
         accumulator.evaluate()?;
-        assert_eq!(
-            Box::new(accumulator).result_owned(),
-            DataValue::Float64(1.5.into())
-        );
+        assert_eq!(accumulator.take_result(), DataValue::Float64(1.5.into()));
         Ok(())
     }
 
@@ -145,16 +152,13 @@ mod tests {
         let mut empty = AvgAccumulator::new(false);
         empty.evaluate()?;
         assert_eq!(empty.result(), &DataValue::Null);
-        assert_eq!(Box::new(empty).result_owned(), DataValue::Null);
+        assert_eq!(empty.take_result(), DataValue::Null);
 
         let mut accumulator = AvgAccumulator::new(false);
         accumulator.update_value(&DataValue::UInt32(2))?;
         accumulator.update_value(&DataValue::UInt32(4))?;
         accumulator.evaluate()?;
-        assert_eq!(
-            Box::new(accumulator).result_owned(),
-            DataValue::Float64(3.0.into())
-        );
+        assert_eq!(accumulator.take_result(), DataValue::Float64(3.0.into()));
         Ok(())
     }
 }