chore: add histogram buckets option on DataBaseBuilder

KKould · KKould · commit 262692de3fa5 · 2026-03-14T15:44:47.000+08:00
diff --git a/src/binder/analyze.rs b/src/binder/analyze.rs
@@ -45,6 +45,7 @@ impl<T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<'_, '_, T, A>
             Operator::Analyze(AnalyzeOperator {
                 table_name,
                 index_metas,
+                histogram_buckets: None,
             }),
             Childrens::Only(Box::new(scan_op)),
         ))
diff --git a/src/db.rs b/src/db.rs
@@ -30,6 +30,7 @@ use crate::optimizer::heuristic::optimizer::HepOptimizerPipeline;
 use crate::optimizer::rule::implementation::ImplementationRuleImpl;
 use crate::optimizer::rule::normalization::NormalizationRuleImpl;
 use crate::parser::parse_sql;
+use crate::planner::operator::Operator;
 use crate::planner::LogicalPlan;
 use crate::storage::memory::MemoryStorage;
 #[cfg(not(target_arch = "wasm32"))]
@@ -64,6 +65,7 @@ pub struct DataBaseBuilder {
     path: PathBuf,
     scala_functions: ScalaFunctions,
     table_functions: TableFunctions,
+    histogram_buckets: Option<usize>,
 }
 
 impl DataBaseBuilder {
@@ -72,6 +74,7 @@ impl DataBaseBuilder {
             path: path.into(),
             scala_functions: Default::default(),
             table_functions: Default::default(),
+            histogram_buckets: None,
         };
         builder = builder.register_scala_function(CharLength::new("char_length".to_lowercase()));
         builder =
@@ -85,6 +88,11 @@ impl DataBaseBuilder {
         builder
     }
 
+    pub fn histogram_buckets(mut self, buckets: usize) -> Self {
+        self.histogram_buckets = Some(buckets);
+        self
+    }
+
     pub fn register_scala_function(mut self, function: Arc<dyn ScalarFunctionImpl>) -> Self {
         let summary = function.summary().clone();
 
@@ -100,41 +108,72 @@ impl DataBaseBuilder {
     }
 
     pub fn build_with_storage<T: Storage>(self, storage: T) -> Result<Database<T>, DatabaseError> {
-        Self::_build::<T>(storage, self.scala_functions, self.table_functions)
+        Self::_build::<T>(
+            storage,
+            self.scala_functions,
+            self.table_functions,
+            self.histogram_buckets,
+        )
     }
 
     #[cfg(target_arch = "wasm32")]
     pub fn build(self) -> Result<Database<MemoryStorage>, DatabaseError> {
         let storage = MemoryStorage::new();
 
-        Self::_build::<MemoryStorage>(storage, self.scala_functions, self.table_functions)
+        Self::_build::<MemoryStorage>(
+            storage,
+            self.scala_functions,
+            self.table_functions,
+            self.histogram_buckets,
+        )
     }
 
     #[cfg(not(target_arch = "wasm32"))]
     pub fn build(self) -> Result<Database<RocksStorage>, DatabaseError> {
         let storage = RocksStorage::new(self.path)?;
 
-        Self::_build::<RocksStorage>(storage, self.scala_functions, self.table_functions)
+        Self::_build::<RocksStorage>(
+            storage,
+            self.scala_functions,
+            self.table_functions,
+            self.histogram_buckets,
+        )
     }
 
     pub fn build_in_memory(self) -> Result<Database<MemoryStorage>, DatabaseError> {
         let storage = MemoryStorage::new();
 
-        Self::_build::<MemoryStorage>(storage, self.scala_functions, self.table_functions)
+        Self::_build::<MemoryStorage>(
+            storage,
+            self.scala_functions,
+            self.table_functions,
+            self.histogram_buckets,
+        )
     }
 
     #[cfg(not(target_arch = "wasm32"))]
     pub fn build_optimistic(self) -> Result<Database<OptimisticRocksStorage>, DatabaseError> {
         let storage = OptimisticRocksStorage::new(self.path)?;
 
-        Self::_build::<OptimisticRocksStorage>(storage, self.scala_functions, self.table_functions)
+        Self::_build::<OptimisticRocksStorage>(
+            storage,
+            self.scala_functions,
+            self.table_functions,
+            self.histogram_buckets,
+        )
     }
 
     fn _build<T: Storage>(
         storage: T,
         scala_functions: ScalaFunctions,
         table_functions: TableFunctions,
+        histogram_buckets: Option<usize>,
     ) -> Result<Database<T>, DatabaseError> {
+        if matches!(histogram_buckets, Some(0)) {
+            return Err(DatabaseError::InvalidValue(
+                "histogram buckets must be >= 1".to_string(),
+            ));
+        }
         let meta_cache = SharedLruCache::new(256, 8, RandomState::new())?;
         let table_cache = SharedLruCache::new(48, 4, RandomState::new())?;
         let view_cache = SharedLruCache::new(12, 4, RandomState::new())?;
@@ -149,6 +188,7 @@ impl DataBaseBuilder {
                 table_cache,
                 view_cache,
                 optimizer_pipeline: default_optimizer_pipeline(),
+                histogram_buckets,
                 _p: Default::default(),
             }),
         })
@@ -260,6 +300,7 @@ pub(crate) struct State<S> {
     table_cache: TableCache,
     view_cache: ViewCache,
     optimizer_pipeline: HepOptimizerPipeline,
+    histogram_buckets: Option<usize>,
     _p: PhantomData<S>,
 }
 
@@ -312,11 +353,17 @@ impl<S: Storage> State<S> {
         ///     Limit(1)
         ///       Project(a,b)
         let source_plan = binder.bind(stmt)?;
-        let best_plan = self
+        let mut best_plan = self
             .optimizer_pipeline
             .instantiate(source_plan)
             .find_best(Some(&transaction.meta_loader(meta_cache)))?;
 
+        if let Operator::Analyze(op) = &mut best_plan.operator {
+            if op.histogram_buckets.is_none() {
+                op.histogram_buckets = self.histogram_buckets;
+            }
+        }
+
         Ok(best_plan)
     }
 
@@ -1036,4 +1083,17 @@ pub(crate) mod test {
 
         Ok(())
     }
+
+    #[test]
+    fn test_invalid_histogram_buckets() {
+        let temp_dir = TempDir::new().expect("unable to create temporary working directory");
+        let result = DataBaseBuilder::path(temp_dir.path())
+            .histogram_buckets(0)
+            .build();
+
+        assert!(matches!(
+            result,
+            Err(DatabaseError::InvalidValue(message)) if message == "histogram buckets must be >= 1"
+        ));
+    }
 }
diff --git a/src/execution/dml/analyze.rs b/src/execution/dml/analyze.rs
@@ -48,6 +48,7 @@ pub struct Analyze {
     table_name: TableName,
     input: LogicalPlan,
     index_metas: Vec<IndexMetaRef>,
+    histogram_buckets: Option<usize>,
 }
 
 impl From<(AnalyzeOperator, LogicalPlan)> for Analyze {
@@ -56,6 +57,7 @@ impl From<(AnalyzeOperator, LogicalPlan)> for Analyze {
             AnalyzeOperator {
                 table_name,
                 index_metas,
+                histogram_buckets,
             },
             input,
         ): (AnalyzeOperator, LogicalPlan),
@@ -64,6 +66,7 @@ impl From<(AnalyzeOperator, LogicalPlan)> for Analyze {
             table_name,
             input,
             index_metas,
+            histogram_buckets,
         }
     }
 }
@@ -79,6 +82,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Analyze {
                 table_name,
                 mut input,
                 index_metas,
+                histogram_buckets,
             } = self;
 
             let schema = input.output_schema().clone();
@@ -99,6 +103,7 @@ impl<'a, T: Transaction + 'a> WriteExecutor<'a, T> for Analyze {
                     index_id: index.id,
                     exprs: throw!(co, index.column_exprs(&table)),
                     builder: HistogramBuilder::new(index, None),
+                    histogram_buckets,
                 });
             }
 
@@ -160,6 +165,7 @@ struct State {
     index_id: IndexId,
     exprs: Vec<ScalarExpression>,
     builder: HistogramBuilder,
+    histogram_buckets: Option<usize>,
 }
 
 impl Analyze {
@@ -177,15 +183,19 @@ impl Analyze {
         let mut active_index_paths = HashSet::new();
 
         for State {
-            index_id, builder, ..
+            index_id,
+            builder,
+            histogram_buckets,
+            ..
         } in builders
         {
             let index_file = OsStr::new(&index_id.to_string()).to_os_string();
             let path = dir_path.join(&index_file);
             let temp_path = path.with_extension("tmp");
             let path_str: String = path.to_string_lossy().into();
 
-            let (histogram, sketch) = builder.build(DEFAULT_NUM_OF_BUCKETS)?;
+            let (histogram, sketch) =
+                builder.build(histogram_buckets.unwrap_or(DEFAULT_NUM_OF_BUCKETS))?;
             let meta = StatisticsMeta::new(histogram, sketch);
 
             meta.to_file(&temp_path)?;
@@ -224,11 +234,15 @@ impl Analyze {
         let mut active_keys = HashSet::new();
 
         for State {
-            index_id, builder, ..
+            index_id,
+            builder,
+            histogram_buckets,
+            ..
         } in builders
         {
             let key = format!("{prefix}/{index_id}");
-            let (histogram, sketch) = builder.build(DEFAULT_NUM_OF_BUCKETS)?;
+            let (histogram, sketch) =
+                builder.build(histogram_buckets.unwrap_or(DEFAULT_NUM_OF_BUCKETS))?;
             let meta = StatisticsMeta::new(histogram, sketch);
             let encoded = meta.to_storage_string()?;
 
@@ -300,7 +314,10 @@ mod test {
     fn test_statistics_meta() -> Result<(), DatabaseError> {
         let temp_dir = TempDir::new().expect("unable to create temporary working directory");
         let base_dir = require_statistics_base_dir();
-        let kite_sql = DataBaseBuilder::path(temp_dir.path()).build()?;
+        let buckets = 10;
+        let kite_sql = DataBaseBuilder::path(temp_dir.path())
+            .histogram_buckets(buckets)
+            .build()?;
 
         kite_sql
             .run("create table t1 (a int primary key, b int)")?
@@ -328,16 +345,19 @@ mod test {
 
         assert_eq!(statistics_meta_pk_index.index_id(), 0);
         assert_eq!(statistics_meta_pk_index.histogram().values_len(), 101);
+        assert_eq!(statistics_meta_pk_index.histogram().buckets_len(), buckets);
 
         let statistics_meta_b_index = StatisticsMeta::from_file::<RocksTransaction>(&paths[1])?;
 
         assert_eq!(statistics_meta_b_index.index_id(), 1);
         assert_eq!(statistics_meta_b_index.histogram().values_len(), 101);
+        assert_eq!(statistics_meta_b_index.histogram().buckets_len(), buckets);
 
         let statistics_meta_p_index = StatisticsMeta::from_file::<RocksTransaction>(&paths[2])?;
 
         assert_eq!(statistics_meta_p_index.index_id(), 2);
         assert_eq!(statistics_meta_p_index.histogram().values_len(), 101);
+        assert_eq!(statistics_meta_p_index.histogram().buckets_len(), buckets);
 
         Ok(())
     }
diff --git a/src/planner/operator/analyze.rs b/src/planner/operator/analyze.rs
@@ -20,4 +20,5 @@ use kite_sql_serde_macros::ReferenceSerialization;
 pub struct AnalyzeOperator {
     pub table_name: TableName,
     pub index_metas: Vec<IndexMetaRef>,
+    pub histogram_buckets: Option<usize>,
 }

Original file line number	Diff line number	Diff line change
`@@ -20,4 +20,5 @@ use kite_sql_serde_macros::ReferenceSerialization;`
`20`	`20`	`pub struct AnalyzeOperator {`
`21`	`21`	`pub table_name: TableName,`
`22`	`22`	`pub index_metas: Vec<IndexMetaRef>,`
	`23`	`+ pub histogram_buckets: Option<usize>,`
`23`	`24`	`}`