feat: add single column compaction for database tool compact command (#…

…9453)

Currently we only support compacting all columns at once. Sometimes it can be useful to compact a single column, which should be much faster. This PR adds support for that.

core/store/src/db/rocksdb.rs

@@ -297,6 +297,13 @@ impl RocksDB {
             _ => unreachable!(),
         }
     }
+
+    pub fn compact_column(&self, col: DBCol) -> io::Result<()> {
+        let none = Option::<&[u8]>::None;
+        tracing::info!(target: "db", column = %col, "Compact column");
+        self.db.compact_range_cf(self.cf_handle(col)?, none, none);
+        Ok(())
+    }
 }
 
 impl Database for RocksDB {
@@ -375,10 +382,8 @@ impl Database for RocksDB {
     }
 
     fn compact(&self) -> io::Result<()> {
-        let none = Option::<&[u8]>::None;
         for col in DBCol::iter() {
-            tracing::info!(target: "db", column = %col, "Compacted column");
-            self.db.compact_range_cf(self.cf_handle(col)?, none, none);
+            self.compact_column(col)?;
         }
         Ok(())
     }

tools/database/src/analyse_data_size_distribution.rs

@@ -8,7 +8,7 @@ use std::sync::{Arc, Mutex};
 use std::{panic, println};
 use strum::IntoEnumIterator;
 
-use crate::utils::open_rocksdb;
+use crate::utils::{open_rocksdb, resolve_column};
 
 #[derive(Parser)]
 pub(crate) struct AnalyseDataSizeDistributionCommand {
@@ -21,10 +21,6 @@ pub(crate) struct AnalyseDataSizeDistributionCommand {
     top_k: usize,
 }
 
-fn resolve_db_col(col: &str) -> Option<DBCol> {
-    DBCol::iter().filter(|db_col| <&str>::from(db_col) == col).next()
-}
-
 #[derive(Clone)]
 struct ColumnFamilyCountAndSize {
     number_of_pairs: usize,
@@ -186,19 +182,17 @@ fn read_all_pairs(db: &RocksDB, col_families: &Vec<DBCol>) -> DataSizeDistributi
     DataSizeDistribution::new(key_sizes, value_sizes, column_families)
 }
 
-fn get_column_families(input_col: &Option<String>) -> Vec<DBCol> {
+fn get_column_families(input_col: &Option<String>) -> anyhow::Result<Vec<DBCol>> {
     match input_col {
-        Some(column_name) => {
-            vec![resolve_db_col(&column_name).unwrap()]
-        }
-        None => DBCol::iter().collect(),
+        Some(column_name) => Ok(vec![resolve_column(column_name)?]),
+        None => Ok(DBCol::iter().collect()),
     }
 }
 
 impl AnalyseDataSizeDistributionCommand {
     pub(crate) fn run(&self, home: &PathBuf) -> anyhow::Result<()> {
         let db = open_rocksdb(home, near_store::Mode::ReadOnly)?;
-        let column_families = get_column_families(&self.column);
+        let column_families = get_column_families(&self.column)?;
         let results = read_all_pairs(&db, &column_families);
         results.print_results(self.top_k);
 

tools/database/src/compact.rs

@@ -1,15 +1,23 @@
-use crate::utils::open_rocksdb;
+use crate::utils::{open_rocksdb, resolve_column};
 use clap::Parser;
 use near_store::db::Database;
 use std::path::PathBuf;
 
 #[derive(Parser)]
-pub(crate) struct RunCompactionCommand {}
+pub(crate) struct RunCompactionCommand {
+    /// If specified only this column will compacted
+    #[arg(short, long)]
+    column: Option<String>,
+}
 
 impl RunCompactionCommand {
     pub(crate) fn run(&self, home: &PathBuf) -> anyhow::Result<()> {
         let db = open_rocksdb(home, near_store::Mode::ReadWrite)?;
-        db.compact()?;
+        if let Some(col_name) = &self.column {
+            db.compact_column(resolve_column(col_name)?)?;
+        } else {
+            db.compact()?;
+        }
         eprintln!("Compaction is finished!");
         Ok(())
     }

tools/database/src/utils.rs

@@ -1,5 +1,9 @@
 use std::path::Path;
 
+use anyhow::anyhow;
+use near_store::DBCol;
+use strum::IntoEnumIterator;
+
 pub(crate) fn open_rocksdb(
     home: &Path,
     mode: near_store::Mode,
@@ -13,3 +17,10 @@ pub(crate) fn open_rocksdb(
         near_store::db::RocksDB::open(&db_path, store_config, mode, near_store::Temperature::Hot)?;
     Ok(rocksdb)
 }
+
+pub(crate) fn resolve_column(col_name: &str) -> anyhow::Result<DBCol> {
+    DBCol::iter()
+        .filter(|db_col| <&str>::from(db_col) == col_name)
+        .next()
+        .ok_or_else(|| anyhow!("column {col_name} does not exist"))
+}