Repository: paritytech/parity-db (language: rust)

Issue:
Expose column stats
Current stats can be written in text form, but it would be useful to access them programmatically too. For instance, I'm interested in number of elements stored in a hash column and/or amount of data.

The data is already there, but no API to access it unfortunately.


Discussion / hints:

Fix patch:
diff --git a/admin/src/lib.rs b/admin/src/lib.rs
index 3512f516..1c43bea8 100644
--- a/admin/src/lib.rs
+++ b/admin/src/lib.rs
@@ -41,7 +41,7 @@ pub fn run() -> Result<(), String> {
 		.unwrap_or_else(|| std::env::current_dir().expect("Cannot resolve current dir"));
 	let nb_column = cli.shared().columns.unwrap_or(1);
 	let mut options = if let Some(metadata) = parity_db::Options::load_metadata(&db_path)
-		.map_err(|e| format!("Error resolving metas: {:?}", e))?
+		.map_err(|e| format!("Error resolving metadata: {:?}", e))?
 	{
 		let mut options = parity_db::Options::with_columns(db_path.as_path(), 0);
 		options.columns = metadata.columns;
@@ -68,7 +68,7 @@ pub fn run() -> Result<(), String> {
 				db.clear_stats(stat.column);
 			} else {
 				let mut out = std::io::stdout();
-				db.collect_stats(&mut out, stat.column);
+				db.write_stats_text(&mut out, stat.column).unwrap();
 			}
 		},
 		SubCommand::Migrate(args) => {
diff --git a/src/column.rs b/src/column.rs
index df60c4c2..a585686e 100644
--- a/src/column.rs
+++ b/src/column.rs
@@ -23,7 +23,7 @@ use crate::{
 	index::{Address, IndexTable, PlanOutcome, TableId as IndexTableId},
 	log::{Log, LogAction, LogOverlays, LogQuery, LogReader, LogWriter},
 	options::{ColumnOptions, Metadata, Options},
-	stats::ColumnStats,
+	stats::{ColumnStatSummary, ColumnStats},
 	table::{
 		key::{TableKey, TableKeyQuery},
 		TableId as ValueTableId, Value, ValueTable, SIZE_TIERS,
@@ -645,16 +645,20 @@ impl HashColumn {
 		Ok(())
 	}
 
-	pub fn write_stats(&self, writer: &mut impl std::io::Write) {
+	pub fn write_stats_text(&self, writer: &mut impl std::io::Write) -> Result<()> {
 		let tables = self.tables.read();
 		tables.index.write_stats(&self.stats);
-		self.stats.write_summary(writer, tables.index.id.col());
+		self.stats.write_stats_text(writer, tables.index.id.col())
 	}
 
-	pub fn clear_stats(&self) {
+	fn stat_summary(&self) -> ColumnStatSummary {
+		self.stats.summary()
+	}
+
+	fn clear_stats(&self) {
 		let tables = self.tables.read();
-		let empty_stats = ColumnStats::empty();
-		tables.index.write_stats(&empty_stats);
+		self.stats.clear();
+		tables.index.write_stats(&self.stats);
 	}
 
 	pub fn iter_while(&self, log: &Log, mut f: impl FnMut(IterState) -> bool) -> Result<()> {
@@ -1031,10 +1035,10 @@ impl Column {
 		}
 	}
 
-	pub fn write_stats(&self, writer: &mut impl std::io::Write) {
+	pub fn write_stats_text(&self, writer: &mut impl std::io::Write) -> Result<()> {
 		match self {
-			Column::Hash(column) => column.write_stats(writer),
-			Column::Tree(_column) => (),
+			Column::Hash(column) => column.write_stats_text(writer),
+			Column::Tree(_column) => Ok(()),
 		}
 	}
 
@@ -1045,6 +1049,13 @@ impl Column {
 		}
 	}
 
+	pub fn stats(&self) -> Option<ColumnStatSummary> {
+		match self {
+			Column::Hash(column) => Some(column.stat_summary()),
+			Column::Tree(_column) => None,
+		}
+	}
+
 	pub fn dump(&self, log: &Log, check_param: &crate::CheckOptions, col: ColId) -> Result<()> {
 		match self {
 			Column::Hash(column) => column.dump(log, check_param, col),
diff --git a/src/db.rs b/src/db.rs
index 279e9949..fee647e7 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -21,6 +21,7 @@ use crate::{
 	index::PlanOutcome,
 	log::{Log, LogAction},
 	options::Options,
+	stats::StatSummary,
 	Key,
 };
 use fs2::FileExt;
@@ -731,7 +732,9 @@ impl DbInner {
 			match std::fs::File::create(path) {
 				Ok(file) => {
 					let mut writer = std::io::BufWriter::new(file);
-					self.collect_stats(&mut writer, None)
+					if let Err(e) = self.write_stats_text(&mut writer, None) {
+						log::warn!(target: "parity-db", "Error writing stats file: {:?}", e)
+					}
 				},
 				Err(e) => log::warn!(target: "parity-db", "Error creating stats file: {:?}", e),
 			}
@@ -739,13 +742,14 @@ impl DbInner {
 		Ok(())
 	}
 
-	fn collect_stats(&self, writer: &mut impl std::io::Write, column: Option<u8>) {
+	fn write_stats_text(&self, writer: &mut impl std::io::Write, column: Option<u8>) -> Result<()> {
 		if let Some(col) = column {
-			self.columns[col as usize].write_stats(writer);
+			self.columns[col as usize].write_stats_text(writer)
 		} else {
 			for c in self.columns.iter() {
-				c.write_stats(writer);
+				c.write_stats_text(writer)?;
 			}
+			Ok(())
 		}
 	}
 
@@ -759,6 +763,10 @@ impl DbInner {
 		}
 	}
 
+	fn stats(&self) -> StatSummary {
+		StatSummary { columns: self.columns.iter().map(|c| c.stats()).collect() }
+	}
+
 	fn store_err(&self, result: Result<()>) {
 		if let Err(e) = result {
 			log::warn!(target: "parity-db", "Background worker error: {}", e);
@@ -968,8 +976,12 @@ impl Db {
 		Ok(())
 	}
 
-	pub fn collect_stats(&self, writer: &mut impl std::io::Write, column: Option<u8>) {
-		self.inner.collect_stats(writer, column)
+	pub fn write_stats_text(
+		&self,
+		writer: &mut impl std::io::Write,
+		column: Option<u8>,
+	) -> Result<()> {
+		self.inner.write_stats_text(writer, column)
 	}
 
 	pub fn clear_stats(&self, column: Option<u8>) {
@@ -986,6 +998,10 @@ impl Db {
 		}
 		Ok(())
 	}
+
+	pub fn stats(&self) -> StatSummary {
+		self.inner.stats()
+	}
 }
 
 impl Drop for Db {
@@ -1286,7 +1302,7 @@ mod tests {
 
 	use super::{Db, EnableCommitPipelineStages, InternalOptions, Options};
 	use rand::Rng;
-	use std::collections::BTreeMap;
+	use std::collections::{BTreeMap, HashSet};
 	use tempfile::tempdir;
 
 	#[test]
@@ -1687,8 +1703,8 @@ mod tests {
 
 	fn test_basic(change_set: &[(Vec<u8>, Option<Vec<u8>>)]) {
 		let tmp = tempdir().unwrap();
-		let col_nb = 0u8;
-		let mut options = Options::with_columns(tmp.path(), 5);
+		let col_nb = 1u8;
+		let mut options = Options::with_columns(tmp.path(), 2);
 		options.columns[col_nb as usize].btree_index = true;
 		let db_test = EnableCommitPipelineStages::DbFile;
 		let inner_options =
@@ -1698,10 +1714,28 @@ mod tests {
 		let mut iter = db.iter(col_nb).unwrap();
 		assert_eq!(iter.next().unwrap(), None);
 
+		db.commit(change_set.iter().map(|(k, v)| (0, k.clone(), v.clone()))).unwrap();
 		db.commit(change_set.iter().map(|(k, v)| (col_nb, k.clone(), v.clone())))
 			.unwrap();
 		db_test.run_stages(&db);
 
+		let stats = db.stats();
+		let col_stats = stats.columns.into_iter().next().unwrap().unwrap();
+		let mut keys = HashSet::new();
+		let mut expected_count: u64 = 0;
+		for (k, v) in change_set.iter() {
+			if v.is_some() {
+				if keys.insert(k) {
+					expected_count += 1;
+				}
+			} else {
+				if keys.remove(k) {
+					expected_count -= 1;
+				}
+			}
+		}
+		assert_eq!(col_stats.total_values, expected_count);
+
 		let state: BTreeMap<Vec<u8>, Option<Vec<u8>>> =
 			change_set.iter().map(|(k, v)| (k.clone(), v.clone())).collect();
 		for (key, value) in state.iter() {
diff --git a/src/lib.rs b/src/lib.rs
index 6774b5eb..5edc239d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -34,6 +34,7 @@ pub use db::{check::CheckOptions, Db, Value};
 pub use error::{Error, Result};
 pub use migration::{clear_column, migrate};
 pub use options::{ColumnOptions, Options};
+pub use stats::{ColumnStatSummary, StatSummary};
 
 #[derive(Default)]
 pub struct IdentityKeyHash(u64);
diff --git a/src/stats.rs b/src/stats.rs
index ff80c273..3754be4f 100644
--- a/src/stats.rs
+++ b/src/stats.rs
@@ -48,6 +48,25 @@ pub struct ColumnStats {
 	compression_delta: [AtomicI64; HISTOGRAM_BUCKETS],
 }
 
+/// Database statistics summary.
+pub struct StatSummary {
+	/// Per column statistics.
+	/// Statistics may be available only for some columns.
+	pub columns: Vec<Option<ColumnStatSummary>>,
+}
+
+/// Column statistics summary.
+pub struct ColumnStatSummary {
+	/// Current number of values in the column.
+	pub total_values: u64,
+	/// Total size of (compressed) values in the column. This does not include key size and any
+	/// other overhead.
+	pub total_bytes: u64,
+	/// Total size of values in the column before compression. This does not include key size and
+	/// any other overhead.
+	pub uncompressed_bytes: u64,
+}
+
 fn read_u32(cursor: &mut Cursor<&[u8]>) -> AtomicU32 {
 	let mut buf = [0u8; 4];
 	cursor.read_exact(&mut buf).expect("Incorrect stats buffer");
@@ -168,6 +187,38 @@ impl ColumnStats {
 		}
 	}
 
+	pub fn clear(&self) {
+		// This may break stat consistency, but we don't care too much.
+		for v in &self.value_histogram {
+			v.store(0, Ordering::Relaxed)
+		}
+		for v in &self.query_histogram {
+			v.store(0, Ordering::Relaxed)
+		}
+		self.oversized.store(0, Ordering::Relaxed);
+		self.oversized_bytes.store(0, Ordering::Relaxed);
+		self.total_values.store(0, Ordering::Relaxed);
+		self.total_bytes.store(0, Ordering::Relaxed);
+		self.commits.store(0, Ordering::Relaxed);
+		self.inserted_new.store(0, Ordering::Relaxed);
+		self.inserted_overwrite.store(0, Ordering::Relaxed);
+		self.removed_hit.store(0, Ordering::Relaxed);
+		self.removed_miss.store(0, Ordering::Relaxed);
+		self.queries_miss.store(0, Ordering::Relaxed);
+		self.uncompressed_bytes.store(0, Ordering::Relaxed);
+		for v in &self.compression_delta {
+			v.store(0, Ordering::Relaxed)
+		}
+	}
+
+	pub fn summary(&self) -> ColumnStatSummary {
+		ColumnStatSummary {
+			total_values: self.total_values.load(Ordering::Relaxed),
+			total_bytes: self.total_bytes.load(Ordering::Relaxed),
+			uncompressed_bytes: self.uncompressed_bytes.load(Ordering::Relaxed),
+		}
+	}
+
 	pub fn to_slice(&self, data: &mut [u8]) {
 		let mut cursor = Cursor::new(data);
 		for item in &self.value_histogram {
@@ -192,7 +243,7 @@ impl ColumnStats {
 		}
 	}
 
-	fn write_stats(&self, writer: &mut impl std::io::Write, col: ColId) -> Result<()> {
+	pub fn write_stats_text(&self, writer: &mut impl std::io::Write, col: ColId) -> Result<()> {
 		writeln!(writer, "Column {}", col)?;
 		writeln!(writer, "Total values: {}", self.total_values.load(Ordering::Relaxed))?;
 		writeln!(writer, "Total bytes: {}", self.total_bytes.load(Ordering::Relaxed))?;
@@ -256,10 +307,6 @@ impl ColumnStats {
 		Ok(())
 	}
 
-	pub fn write_summary(&self, writer: &mut impl std::io::Write, col: ColId) {
-		let _ = self.write_stats(writer, col);
-	}
-
 	pub fn query_hit(&self, size_tier: u8) {
 		self.query_histogram[size_tier as usize].fetch_add(1, Ordering::Relaxed);
 	}


Test patch: