rocksdb: trivial count

Signed-off-by: Duc Canh Le <duccanh.le@ahrefs.com>

src/Core/Settings.h

@@ -546,6 +546,7 @@ class IColumn;
     M(Bool, database_atomic_wait_for_drop_and_detach_synchronously, false, "When executing DROP or DETACH TABLE in Atomic database, wait for table data to be finally dropped or detached.", 0) \
     M(Bool, enable_scalar_subquery_optimization, true, "If it is set to true, prevent scalar subqueries from (de)serializing large scalar values and possibly avoid running the same subquery more than once.", 0) \
     M(Bool, optimize_trivial_count_query, true, "Process trivial 'SELECT count() FROM table' query from metadata.", 0) \
+    M(Bool, rocksdb_enable_approximate_count, true, "If `optimize_trivial_count_query` is true, process trivial 'SELECT count() FROM rockdb_table' from metadata, the returned result is approximated.", 0) \
     M(Bool, optimize_count_from_files, true, "Optimize counting rows from files in supported input formats", 0) \
     M(Bool, use_cache_for_count_from_files, true, "Use cache to count the number of rows in files", 0) \
     M(Bool, optimize_respect_aliases, true, "If it is set to true, it will respect aliases in WHERE/GROUP BY/ORDER BY, that will help with partition pruning/secondary indexes/optimize_aggregation_in_order/optimize_read_in_order/optimize_trivial_count", 0) \

src/Storages/RocksDB/StorageEmbeddedRocksDB.cpp

@@ -609,5 +609,19 @@ void registerStorageEmbeddedRocksDB(StorageFactory & factory)
     factory.registerStorage("EmbeddedRocksDB", create, features);
 }
 
+std::optional<UInt64> StorageEmbeddedRocksDB::totalRows(const Settings & settings) const
+{
+    if (settings.rocksdb_enable_approximate_count)
+    {
+        std::shared_lock lock(rocksdb_ptr_mx);
+        if (!rocksdb_ptr)
+            return {};
+        UInt64 estimated_rows;
+        if (!rocksdb_ptr->GetIntProperty("rocksdb.estimate-num-keys", &estimated_rows))
+            return {};
+        return estimated_rows;
+    }
+    return {};
+}
 
 }

src/Storages/RocksDB/StorageEmbeddedRocksDB.h

@@ -83,6 +83,10 @@ public:
 
     bool supportsDelete() const override { return true; }
 
+    bool supportsTrivialCountOptimization() const override { return true; }
+
+    std::optional<UInt64> totalRows(const Settings & settings) const override;
+
 private:
     const String primary_key;
     using RocksDBPtr = std::unique_ptr<rocksdb::DB>;

tests/queries/0_stateless/02892_rocksdb_trivial_count.reference

@@ -0,0 +1,2 @@
+121
+121

tests/queries/0_stateless/02892_rocksdb_trivial_count.sql

@@ -0,0 +1,4 @@
+CREATE TABLE dict (key UInt64, value String) ENGINE = EmbeddedRocksDB PRIMARY KEY key;
+INSERT INTO dict SELECT number, toString(number) FROM numbers(121);
+SELECT count() FROM dict SETTINGS rocksdb_enable_approximate_count = 0;
+SELECT count() FROM dict SETTINGS optimize_trivial_count_query = 1, rocksdb_enable_approximate_count = 1;