You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by go...@apache.org on 2026/09/03 09:11:38 UTC
(pinot) 02/02: Add SQLite and in-memory statistics stores
This is an automated email from the ASF dual-hosted git repository.
gortiz pushed a commit to branch cbo-1-stats-contracts-and-stores
in repository https://gitbox.apache.org/repos/asf/pinot.git
commit b8878b58f943f7dce73f10676051d69155da636c
Author: Gonzalo Ortiz <19...@users.noreply.github.com>
AuthorDate: Thu Aug 27 20:02:17 2026 +0200
Add SQLite and in-memory statistics stores
Statistics for hundreds of thousands of segments must not compete with query
execution for JVM heap, so the default store is an embedded SQLite database
(WAL, a small read-connection pool, and drop-and-rebuild recovery if the file is
ever unreadable). Surviving a restart is a side benefit rather than the goal:
the file also lets the broker skip re-collecting segments whose crc has not
changed.
There is no migration framework. SQLite's own user_version pragma carries the
schema version, and a store whose version does not match is discarded and
rebuilt rather than migrated. Every row here is derived from ZooKeeper metadata
the broker re-reads at startup, so a rebuild costs nothing that was not already
being read -- and it puts schema change, corruption and an unreadable file on
one recovery path. This adds exactly one third-party dependency, sqlite-jdbc.
InMemoryStatsStore is the alternative for brokers that cannot or should not
write a file. It keeps the same semantics -- consuming segments excluded from
aggregates but present for crc reconciliation, no-data reported as absent rather
than zero -- by folding through the same StatsAggregations helpers.
Both serve table-level reads from a rollup recomputed only after a write, since
query planning asks for them on every compile while writes arrive at segment
push cadence. The version stamp is what makes that safe without holding the
write lock: a rollup computed from rows that changed underneath it is used for
that call but never published.
Tests are written as a contract both stores must satisfy -- the optimizer must
not change with the configured store -- with each implementation adding only
what is specific to it: durability and corruption recovery for SQLite, starting
empty for in-memory.
---
LICENSE-binary | 4 +
pinot-broker/pom.xml | 8 +
.../pinot/broker/stats/InMemoryStatsStore.java | 357 +++++++++
.../pinot/broker/stats/SqliteStatsStore.java | 885 +++++++++++++++++++++
.../pinot/broker/stats/InMemoryStatsStoreTest.java | 52 ++
.../pinot/broker/stats/SqliteStatsStoreTest.java | 259 ++++++
.../pinot/broker/stats/StatsStoreContractTest.java | 560 +++++++++++++
pom.xml | 8 +
8 files changed, 2133 insertions(+)
diff --git a/LICENSE-binary b/LICENSE-binary
index 30034937325..956d81c0d6f 100644
--- a/LICENSE-binary
+++ b/LICENSE-binary
@@ -542,6 +542,7 @@ org.scala-lang:scala-reflect:2.13.18
org.slf4j:jcl-over-slf4j:2.0.17
org.webjars:swagger-ui:5.32.0
org.xerial.snappy:snappy-java:1.1.10.8
+org.xerial:sqlite-jdbc:3.46.1.3
org.yaml:snakeyaml:2.6
software.amazon.awssdk:annotations:2.42.11
software.amazon.awssdk:apache-client:2.42.11
@@ -627,6 +628,9 @@ BSD 2-Clause
------------
com.github.luben:zstd-jni:1.5.7-7
org.codehaus.woodstox:stax2-api:4.2.2
+org.xerial:sqlite-jdbc:3.46.1.3 (the JDBC layer derived from Zentus SQLiteJDBC;
+ the remainder of the artifact is Apache-2.0, listed above)
+ Copyright (c) 2006, David Crawshaw. All rights reserved.
BSD 3-Clause
diff --git a/pinot-broker/pom.xml b/pinot-broker/pom.xml
index b96db348c40..ba90112a91b 100644
--- a/pinot-broker/pom.xml
+++ b/pinot-broker/pom.xml
@@ -33,6 +33,10 @@
<pinot.root>${basedir}/..</pinot.root>
</properties>
<dependencies>
+ <dependency>
+ <groupId>org.apache.pinot</groupId>
+ <artifactId>pinot-query-planner-spi</artifactId>
+ </dependency>
<dependency>
<groupId>org.apache.pinot</groupId>
<artifactId>pinot-materialized-view</artifactId>
@@ -53,6 +57,10 @@
<groupId>org.slf4j</groupId>
<artifactId>slf4j-api</artifactId>
</dependency>
+ <dependency>
+ <groupId>org.xerial</groupId>
+ <artifactId>sqlite-jdbc</artifactId>
+ </dependency>
<!-- Test -->
<dependency>
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/stats/InMemoryStatsStore.java b/pinot-broker/src/main/java/org/apache/pinot/broker/stats/InMemoryStatsStore.java
new file mode 100644
index 00000000000..22f7ee8ef11
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/stats/InMemoryStatsStore.java
@@ -0,0 +1,357 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.stats;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicLong;
+import javax.annotation.Nullable;
+import org.apache.pinot.query.planner.spi.stats.ColumnStatistics;
+import org.apache.pinot.query.planner.spi.stats.SegmentColumnStatsRow;
+import org.apache.pinot.query.planner.spi.stats.SegmentStatsRow;
+import org.apache.pinot.query.planner.spi.stats.StatConfidence;
+import org.apache.pinot.query.planner.spi.stats.StatsAggregations;
+import org.apache.pinot.query.planner.spi.stats.StatsStore;
+import org.apache.pinot.query.planner.spi.stats.StatsStoreException;
+import org.apache.pinot.query.planner.spi.stats.TableStatistics;
+
+
+/// Heap-resident [StatsStore], for deployments that do not want a database file on the broker
+/// (read-only filesystems, ephemeral containers, tests).
+///
+/// ### What it trades away
+/// Statistics do not survive a restart. For the base tier that costs little: the values are derived
+/// from the ZooKeeper segment metadata the broker re-reads at startup anyway, so an empty store is
+/// simply re-populated by the same listener callbacks — [#getSegmentCrcs] returns an empty map and
+/// every segment is treated as new. Per-column statistics pulled from servers are the expensive
+/// case, since those must be fetched again.
+///
+/// Table-level reads (row counts, consuming-segment detection) are served from a rollup that is
+/// recomputed only after a write, so query planning does not scan the segment map. Time-range
+/// estimates still scan the table's committed segments, since the answer depends on the requested
+/// range.
+///
+/// Heap use is proportional to the number of segments (and stored columns) of the tables this
+/// broker serves. The SQLite-backed store exists precisely to keep that off-heap, so prefer this
+/// implementation only when the segment count is modest or a file is unacceptable.
+///
+/// ### Semantics
+/// Identical to the SQLite-backed store, and deliberately expressed through the same
+/// [StatsAggregations] helpers so the two cannot drift: consuming segments are excluded from
+/// table-level, column-level and time-range reads but included by [#getSegmentCrcs]; a column row
+/// only counts while its segment row exists; "no data" is reported as `null` /
+/// [OptionalLong#empty()] and never as a zero estimate.
+///
+/// ### Thread-safety
+/// Safe for concurrent reads with a single concurrent writer, as the interface requires. Reads
+/// aggregate over concurrent maps and therefore observe a weakly-consistent view: a read racing a
+/// multi-row upsert may see part of that batch. Unlike a SQLite transaction the batch is not
+/// atomic, which is acceptable because every consumer treats these values as estimates.
+public class InMemoryStatsStore implements StatsStore {
+
+ /// table name → segment name → row.
+ private final Map<String, Map<String, StoredSegment>> _segments = new ConcurrentHashMap<>();
+ /// table name → segment name → column name → row.
+ private final Map<String, Map<String, Map<String, SegmentColumnStatsRow>>> _columns = new ConcurrentHashMap<>();
+ /// table name → monotonically increasing write counter, bumped after every mutation of that table.
+ private final Map<String, AtomicLong> _versions = new ConcurrentHashMap<>();
+ /// table name → table-level rollup, valid only while its version matches the table's counter.
+ private final Map<String, CachedAggregate> _aggregates = new ConcurrentHashMap<>();
+
+ private volatile boolean _closed;
+
+ @Override
+ public void init() {
+ // Nothing to open: the maps are ready on construction.
+ }
+
+ @Override
+ public void upsertSegmentStats(String tableNameWithType, List<SegmentStatsRow> rows)
+ throws StatsStoreException {
+ checkOpen();
+ long now = System.currentTimeMillis();
+ Map<String, StoredSegment> table = _segments.computeIfAbsent(tableNameWithType, k -> new ConcurrentHashMap<>());
+ for (SegmentStatsRow row : rows) {
+ table.put(row.segmentName(), new StoredSegment(row, now));
+ }
+ invalidate(tableNameWithType);
+ }
+
+ @Override
+ public void upsertSegmentColumnStats(String tableNameWithType, List<SegmentColumnStatsRow> rows)
+ throws StatsStoreException {
+ checkOpen();
+ Map<String, Map<String, SegmentColumnStatsRow>> table =
+ _columns.computeIfAbsent(tableNameWithType, k -> new ConcurrentHashMap<>());
+ for (SegmentColumnStatsRow row : rows) {
+ table.computeIfAbsent(row.segmentName(), k -> new ConcurrentHashMap<>()).put(row.columnName(), row);
+ }
+ }
+
+ @Override
+ public void removeSegments(String tableNameWithType, Collection<String> segmentNames)
+ throws StatsStoreException {
+ checkOpen();
+ Map<String, StoredSegment> segments = _segments.get(tableNameWithType);
+ Map<String, Map<String, SegmentColumnStatsRow>> columns = _columns.get(tableNameWithType);
+ for (String segmentName : segmentNames) {
+ if (segments != null) {
+ segments.remove(segmentName);
+ }
+ if (columns != null) {
+ columns.remove(segmentName);
+ }
+ }
+ // Drop emptied tables, so getTables() means "holds statistics" here exactly as it does in a
+ // row-backed store -- it is the input to a destructive purge, so the two must not disagree.
+ if (segments != null && segments.isEmpty()) {
+ _segments.remove(tableNameWithType, segments);
+ }
+ if (columns != null && columns.isEmpty()) {
+ _columns.remove(tableNameWithType, columns);
+ }
+ if (_segments.get(tableNameWithType) == null && _columns.get(tableNameWithType) == null) {
+ // Nothing left to roll up; see purgeTable for why this removes instead of invalidating.
+ _aggregates.remove(tableNameWithType);
+ _versions.remove(tableNameWithType);
+ } else {
+ invalidate(tableNameWithType);
+ }
+ }
+
+ @Override
+ public Map<String, Long> getSegmentCrcs(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ Map<String, StoredSegment> segments = _segments.get(tableNameWithType);
+ // Mutable and detached, like the SQLite store's result: the SPI does not promise immutability
+ // and callers may adjust the map.
+ if (segments == null) {
+ return new HashMap<>();
+ }
+ // Includes consuming segments: reconciliation must not re-upsert them on every restart.
+ Map<String, Long> crcs = new HashMap<>(segments.size());
+ for (Map.Entry<String, StoredSegment> entry : segments.entrySet()) {
+ crcs.put(entry.getKey(), entry.getValue()._row.crc());
+ }
+ return crcs;
+ }
+
+ @Override
+ @Nullable
+ public TableStatistics getTableStats(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ return aggregate(tableNameWithType)._stats;
+ }
+
+ @Override
+ @Nullable
+ public ColumnStatistics getColumnStats(String tableNameWithType, String columnName)
+ throws StatsStoreException {
+ checkOpen();
+ Map<String, Map<String, SegmentColumnStatsRow>> columnsBySegment = _columns.get(tableNameWithType);
+ Map<String, StoredSegment> segments = _segments.get(tableNameWithType);
+ if (columnsBySegment == null || segments == null) {
+ return null;
+ }
+
+ // Same fold as the SQLite store, through the shared accumulator, so the two cannot disagree.
+ StatsAggregations.ColumnStatsAccumulator accumulator = new StatsAggregations.ColumnStatsAccumulator();
+ for (Map.Entry<String, Map<String, SegmentColumnStatsRow>> entry : columnsBySegment.entrySet()) {
+ SegmentColumnStatsRow column = entry.getValue().get(columnName);
+ if (column == null) {
+ continue;
+ }
+ // Mirrors the SQL join onto segment_stats: a column row without a live, committed segment row
+ // contributes nothing, and the segment row supplies the doc count used for weighting.
+ StoredSegment segment = segments.get(entry.getKey());
+ if (segment == null || segment._row.consuming()) {
+ continue;
+ }
+ accumulator.add(segment._row.totalDocs(), column);
+ }
+ return accumulator.isEmpty() ? null : accumulator.build(columnName);
+ }
+
+ @Override
+ public OptionalLong estimateRowsInTimeRange(String tableNameWithType, long startMs, long endMs)
+ throws StatsStoreException {
+ checkOpen();
+ Map<String, StoredSegment> segments = _segments.get(tableNameWithType);
+ if (segments == null) {
+ return OptionalLong.empty();
+ }
+ long totalRows = 0;
+ boolean hasCommittedSegment = false;
+ for (StoredSegment segment : segments.values()) {
+ SegmentStatsRow row = segment._row;
+ if (row.consuming()) {
+ continue;
+ }
+ // Distinguishes "no overlapping segment" (a real estimate of 0) from "no statistics"
+ // (empty), exactly as the SQL existence sentinel does.
+ hasCommittedSegment = true;
+ totalRows += StatsAggregations.overlapRows(row.totalDocs(), row.startTimeMs(), row.endTimeMs(),
+ startMs, endMs);
+ }
+ return hasCommittedSegment ? OptionalLong.of(totalRows) : OptionalLong.empty();
+ }
+
+ @Override
+ public Set<String> getTables()
+ throws StatsStoreException {
+ checkOpen();
+ Set<String> tables = new HashSet<>(_segments.keySet());
+ tables.addAll(_columns.keySet());
+ return tables;
+ }
+
+ @Override
+ public void purgeTable(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ _segments.remove(tableNameWithType);
+ _columns.remove(tableNameWithType);
+ // Remove rather than invalidate: invalidate() would re-create the version entry it just
+ // dropped, so a broker that serves many tables over its lifetime would accumulate one entry
+ // per table it ever saw.
+ _aggregates.remove(tableNameWithType);
+ _versions.remove(tableNameWithType);
+ }
+
+ @Override
+ public void purgeAll()
+ throws StatsStoreException {
+ checkOpen();
+ _segments.clear();
+ _columns.clear();
+ _aggregates.clear();
+ _versions.clear();
+ }
+
+ @Override
+ public boolean hasConsumingSegments(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ return aggregate(tableNameWithType)._hasConsuming;
+ }
+
+ @Override
+ public void close() {
+ _closed = true;
+ _segments.clear();
+ _columns.clear();
+ _aggregates.clear();
+ _versions.clear();
+ }
+
+ /// Returns the table-level rollup, recomputing it only when the table changed since it was last
+ /// computed. Query planning asks for these on every compile, so an O(#segments) scan per call
+ /// would show up directly in planning latency; writes are comparatively rare.
+ ///
+ /// The version is read before the scan and re-checked after it, so a rollup computed from a state
+ /// a concurrent writer has already replaced is used once but never cached.
+ private CachedAggregate aggregate(String tableNameWithType) {
+ AtomicLong version = _versions.computeIfAbsent(tableNameWithType, k -> new AtomicLong());
+ long observed = version.get();
+ CachedAggregate cached = _aggregates.get(tableNameWithType);
+ if (cached != null && cached._version == observed) {
+ return cached;
+ }
+ CachedAggregate computed = computeAggregate(tableNameWithType, observed);
+ if (version.get() == observed) {
+ _aggregates.put(tableNameWithType, computed);
+ }
+ return computed;
+ }
+
+ private CachedAggregate computeAggregate(String tableNameWithType, long version) {
+ Map<String, StoredSegment> segments = _segments.get(tableNameWithType);
+ if (segments == null) {
+ return new CachedAggregate(version, null, false);
+ }
+ long totalDocs = 0;
+ long sizeBytes = 0;
+ long maxUpdatedAt = 0;
+ int committed = 0;
+ boolean hasConsuming = false;
+ for (StoredSegment segment : segments.values()) {
+ if (segment._row.consuming()) {
+ hasConsuming = true;
+ continue;
+ }
+ totalDocs += segment._row.totalDocs();
+ sizeBytes += segment._row.sizeBytes();
+ maxUpdatedAt = Math.max(maxUpdatedAt, segment._updatedAtMs);
+ committed++;
+ }
+ if (committed == 0) {
+ return new CachedAggregate(version, null, hasConsuming);
+ }
+ return new CachedAggregate(version, TableStatistics.builder()
+ .rowCount(totalDocs, StatConfidence.EXACT)
+ .tableSizeBytes(sizeBytes, StatConfidence.EXACT)
+ .updatedAtMs(maxUpdatedAt)
+ .build(), hasConsuming);
+ }
+
+ private void invalidate(String tableNameWithType) {
+ _versions.computeIfAbsent(tableNameWithType, k -> new AtomicLong()).incrementAndGet();
+ }
+
+ private void checkOpen()
+ throws StatsStoreException {
+ if (_closed) {
+ throw new StatsStoreException("InMemoryStatsStore is closed");
+ }
+ }
+
+ /// A table-level rollup together with the write version it was computed from.
+ private static final class CachedAggregate {
+ private final long _version;
+ @Nullable
+ private final TableStatistics _stats;
+ private final boolean _hasConsuming;
+
+ CachedAggregate(long version, @Nullable TableStatistics stats, boolean hasConsuming) {
+ _version = version;
+ _stats = stats;
+ _hasConsuming = hasConsuming;
+ }
+ }
+
+ /// A stored segment row plus the wall-clock time it was written, which the SQLite schema keeps in
+ /// its `updated_at_ms` column and exposes as [TableStatistics#getUpdatedAtMs()].
+ private static final class StoredSegment {
+ private final SegmentStatsRow _row;
+ private final long _updatedAtMs;
+
+ StoredSegment(SegmentStatsRow row, long updatedAtMs) {
+ _row = row;
+ _updatedAtMs = updatedAtMs;
+ }
+ }
+}
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/stats/SqliteStatsStore.java b/pinot-broker/src/main/java/org/apache/pinot/broker/stats/SqliteStatsStore.java
new file mode 100644
index 00000000000..2b61c6370ad
--- /dev/null
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/stats/SqliteStatsStore.java
@@ -0,0 +1,885 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.stats;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.Set;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
+import javax.annotation.Nullable;
+import org.apache.pinot.query.planner.spi.stats.ColumnStatistics;
+import org.apache.pinot.query.planner.spi.stats.ColumnValueType;
+import org.apache.pinot.query.planner.spi.stats.SegmentColumnStatsRow;
+import org.apache.pinot.query.planner.spi.stats.SegmentStatsRow;
+import org.apache.pinot.query.planner.spi.stats.StatConfidence;
+import org.apache.pinot.query.planner.spi.stats.StatsAggregations;
+import org.apache.pinot.query.planner.spi.stats.StatsStore;
+import org.apache.pinot.query.planner.spi.stats.StatsStoreException;
+import org.apache.pinot.query.planner.spi.stats.TableStatistics;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/// SQLite-backed implementation of [StatsStore].
+///
+/// ### Threading model
+/// Uses a single shared writer [Connection] guarded by `synchronized` on
+/// `_writeLock`, plus a small pool (the configured pool size) of read-only connections
+/// served via a blocking queue. Reads from multiple threads proceed concurrently (each
+/// borrows a connection from the pool, uses it, then returns it). The writer connection sets
+/// `PRAGMA journal_mode=WAL` so readers never block writers.
+///
+/// ### Corruption handling
+/// [#init()] attempts to open the database and apply its schema. On any failure it logs a
+/// warning, deletes the DB file and its WAL/SHM siblings, then retries once from scratch.
+/// Only a second consecutive failure is propagated as [StatsStoreException].
+public class SqliteStatsStore implements StatsStore {
+ private static final Logger LOGGER = LoggerFactory.getLogger(SqliteStatsStore.class);
+
+ /// Default number of pooled readers when the operator sets none.
+ static final int DEFAULT_READ_POOL_SIZE = 4;
+
+ /// Default wait for a pooled reader. Deliberately short: this sits on the query-planning path,
+ /// where a late estimate is worth less than a fast plan without one, and the caller degrades to
+ /// heuristics rather than failing.
+ static final long DEFAULT_READ_BORROW_TIMEOUT_MS = 50;
+
+ private static final String DB_FILE_NAME = "broker-stats.sqlite";
+
+ /// Bump when [#SCHEMA_DDL] changes. Note that the persisted vocabulary is wider than the DDL:
+ /// `value_type` stores [ColumnValueType] names, which a newer build may extend without touching
+ /// the DDL. That case does not need a bump, because an unrecognized name resolves to `null` and
+ /// degrades to untrusted bounds rather than failing the read.
+ ///
+ /// There is no migration path on purpose: every row here is
+ /// derived from the ZooKeeper segment metadata the broker re-reads at startup, so a store whose
+ /// schema does not match is discarded and rebuilt rather than migrated. That keeps schema
+ /// changes, corruption and an unreadable file on one recovery path.
+ private static final int SCHEMA_VERSION = 2;
+
+ private static final String[] SCHEMA_DDL = {
+ "CREATE TABLE segment_stats ("
+ + "table_name TEXT NOT NULL, segment_name TEXT NOT NULL, crc INTEGER NOT NULL, "
+ + "total_docs INTEGER NOT NULL, size_bytes INTEGER NOT NULL, start_time_ms INTEGER NOT NULL, "
+ + "end_time_ms INTEGER NOT NULL, consuming INTEGER NOT NULL, updated_at_ms INTEGER NOT NULL, "
+ + "PRIMARY KEY (table_name, segment_name))",
+ "CREATE TABLE segment_col_stats ("
+ + "table_name TEXT NOT NULL, segment_name TEXT NOT NULL, column_name TEXT NOT NULL, "
+ + "ndv INTEGER NOT NULL, min_value TEXT, max_value TEXT, min_trusted INTEGER NOT NULL, "
+ + "avg_bytes REAL NOT NULL, null_fraction REAL NOT NULL, value_type TEXT, "
+ + "updated_at_ms INTEGER NOT NULL, "
+ + "PRIMARY KEY (table_name, segment_name, column_name))",
+ // getColumnStats filters by (table_name, column_name); column_name is the third key component
+ // so the primary key cannot serve it.
+ "CREATE INDEX idx_segment_col_stats_col ON segment_col_stats(table_name, column_name)",
+ "CREATE INDEX idx_segment_stats_time ON segment_stats(table_name, start_time_ms, end_time_ms)",
+ };
+
+ // SQL constants
+ private static final String SQL_UPSERT_SEGMENT =
+ "INSERT INTO segment_stats(table_name,segment_name,crc,total_docs,size_bytes,"
+ + "start_time_ms,end_time_ms,consuming,updated_at_ms) VALUES(?,?,?,?,?,?,?,?,?) "
+ + "ON CONFLICT(table_name,segment_name) DO UPDATE SET "
+ + "crc=excluded.crc,total_docs=excluded.total_docs,size_bytes=excluded.size_bytes,"
+ + "start_time_ms=excluded.start_time_ms,end_time_ms=excluded.end_time_ms,"
+ + "consuming=excluded.consuming,updated_at_ms=excluded.updated_at_ms";
+
+ private static final String SQL_UPSERT_COL =
+ "INSERT INTO segment_col_stats(table_name,segment_name,column_name,ndv,min_value,"
+ + "max_value,min_trusted,avg_bytes,null_fraction,value_type,updated_at_ms) "
+ + "VALUES(?,?,?,?,?,?,?,?,?,?,?) "
+ + "ON CONFLICT(table_name,segment_name,column_name) DO UPDATE SET "
+ + "ndv=excluded.ndv,min_value=excluded.min_value,max_value=excluded.max_value,"
+ + "min_trusted=excluded.min_trusted,avg_bytes=excluded.avg_bytes,"
+ + "null_fraction=excluded.null_fraction,value_type=excluded.value_type,"
+ + "updated_at_ms=excluded.updated_at_ms";
+
+ private static final String SQL_DELETE_SEGMENT =
+ "DELETE FROM segment_stats WHERE table_name=? AND segment_name=?";
+
+ private static final String SQL_DELETE_COL =
+ "DELETE FROM segment_col_stats WHERE table_name=? AND segment_name=?";
+
+ private static final String SQL_GET_CRCS =
+ "SELECT segment_name,crc FROM segment_stats WHERE table_name=?";
+
+ /// Rollup and the consuming flag in one scan: both are asked for on every plan, and splitting
+ /// them into two statements walked the same rows twice.
+ private static final String SQL_TABLE_STATS =
+ "SELECT SUM(CASE WHEN consuming=0 THEN total_docs ELSE 0 END),"
+ + "SUM(CASE WHEN consuming=0 THEN size_bytes ELSE 0 END),"
+ + "MAX(CASE WHEN consuming=0 THEN updated_at_ms ELSE 0 END),"
+ + "SUM(CASE WHEN consuming=0 THEN 1 ELSE 0 END),"
+ + "MAX(consuming) FROM segment_stats WHERE table_name=?";
+
+ private static final String SQL_COL_STATS =
+ "SELECT s.total_docs,c.ndv,c.min_value,c.max_value,c.min_trusted,c.avg_bytes,c.null_fraction,"
+ + "c.value_type "
+ + "FROM segment_col_stats c "
+ + "JOIN segment_stats s ON s.table_name=c.table_name AND s.segment_name=c.segment_name "
+ + "WHERE c.table_name=? AND c.column_name=? AND s.consuming=0";
+
+ // The overlap predicate mirrors the Java-side check in estimateRowsInTimeRange: it prunes
+ // segments that cannot overlap [startMs, endMs) inside SQLite instead of materializing them
+ // over JDBC only to be skipped. Segments with unknown times (-1 sentinels) must be retained —
+ // they are included conservatively by the Java logic.
+ //
+ // The UNION ALL arm emits a sentinel row (start_time_ms = -2, an otherwise impossible value)
+ // whenever the table has any committed stats at all. It lets the caller distinguish "no segment
+ // overlaps the range" (a real estimate of 0) from "no stats for this table" (empty) within a
+ // SINGLE statement — i.e. a single consistent snapshot; splitting this into a second existence
+ // query would race with concurrent stats writes.
+ private static final String SQL_TIME_RANGE =
+ "SELECT total_docs,start_time_ms,end_time_ms "
+ + "FROM segment_stats WHERE table_name=? AND consuming=0 "
+ + "AND (start_time_ms=-1 OR end_time_ms=-1 OR (end_time_ms>? AND start_time_ms<?)) "
+ + "UNION ALL SELECT -1,-2,-2 WHERE EXISTS("
+ + "SELECT 1 FROM segment_stats WHERE table_name=? AND consuming=0)";
+
+ private static final String SQL_HAS_CONSUMING =
+ "SELECT 1 FROM segment_stats WHERE table_name=? AND consuming=1 LIMIT 1";
+
+ // Both tables are consulted: a table could in principle have column rows without segment rows.
+ private static final String SQL_LIST_TABLES =
+ "SELECT table_name FROM segment_stats UNION SELECT table_name FROM segment_col_stats";
+
+ private static final String SQL_PURGE_TABLE_SEG =
+ "DELETE FROM segment_stats WHERE table_name=?";
+
+ private static final String SQL_PURGE_TABLE_COL =
+ "DELETE FROM segment_col_stats WHERE table_name=?";
+
+ private static final String SQL_PURGE_ALL_SEG = "DELETE FROM segment_stats";
+ private static final String SQL_PURGE_ALL_COL = "DELETE FROM segment_col_stats";
+
+ private final Path _dbDirectory;
+ private final Path _dbPath;
+
+ private final int _readPoolSize;
+ private final long _readBorrowTimeoutMs;
+
+ /// Write connection — all mutations go through this; guarded by _writeLock.
+ /// Volatile because [#openConnections()] and [#closeConnections()] assign it outside the
+ /// `_writeLock` that guards every other access, during init and rebuild.
+ private volatile Connection _writeConn;
+ /// Serialises writers. Needed even though SQLite is itself thread-safe, because a write here is
+ /// a multi-statement transaction (`addBatch` / `executeBatch` / `commit`) and the transaction
+ /// boundary is per-CONNECTION state: two threads interleaving on this shared connection would
+ /// let one thread's `commit()` publish another's half-written batch.
+ ///
+ /// It does NOT make readers wait. The database runs in WAL mode, so a writer never blocks a
+ /// reader and a reader never blocks the writer; this lock is only ever contended between
+ /// concurrent writers.
+ private final Object _writeLock = new Object();
+
+ /// Pool of read-only connections. Each reader borrows a connection, uses it,
+ /// then returns it via `offer()`. Sized at the configured pool size.
+ private final ArrayBlockingQueue<PooledReader> _readPool;
+
+ /// Per-table write counter backing the rollup cache; see [#aggregate(String)].
+ private final Map<String, AtomicLong> _versions = new ConcurrentHashMap<>();
+
+ /// Cached per-table rollups, each stamped with the [#_versions] value it was computed from.
+ private final Map<String, CachedAggregate> _aggregates = new ConcurrentHashMap<>();
+
+ private volatile boolean _closed = false;
+
+ /// Constructs a new `SqliteStatsStore` that stores its database in the given directory.
+ /// The database file will be `<dbDirectory>/broker-stats.sqlite`.
+ ///
+ /// @param dbDirectory directory in which to store the database file; created if absent
+ public SqliteStatsStore(Path dbDirectory) {
+ this(dbDirectory, DEFAULT_READ_POOL_SIZE, DEFAULT_READ_BORROW_TIMEOUT_MS);
+ }
+
+ /// @param dbDirectory directory in which to store the database file; created if absent
+ /// @param readPoolSize number of pooled readers; the pool is a bound, so this caps how
+ /// many estimates can be served concurrently
+ /// @param readBorrowTimeoutMs how long a reader waits for a pooled connection before giving up
+ /// and reporting no statistics
+ public SqliteStatsStore(Path dbDirectory, int readPoolSize, long readBorrowTimeoutMs) {
+ _dbDirectory = dbDirectory;
+ _dbPath = dbDirectory.resolve(DB_FILE_NAME);
+ _readPoolSize = readPoolSize > 0 ? readPoolSize : DEFAULT_READ_POOL_SIZE;
+ _readBorrowTimeoutMs = readBorrowTimeoutMs > 0 ? readBorrowTimeoutMs : DEFAULT_READ_BORROW_TIMEOUT_MS;
+ _readPool = new ArrayBlockingQueue<>(_readPoolSize);
+ }
+
+ // ---------------------------------------------------------------------------
+ // Lifecycle
+ // ---------------------------------------------------------------------------
+
+ @Override
+ public void init()
+ throws StatsStoreException {
+ try {
+ openAndApplySchema();
+ } catch (Exception firstEx) {
+ LOGGER.warn(
+ "Failed to open stats store at {}; deleting and retrying from scratch. Cause: {}",
+ _dbPath, firstEx.getMessage(), firstEx);
+ closeConnections();
+ deleteDbFiles();
+ try {
+ openAndApplySchema();
+ } catch (Exception secondEx) {
+ throw new StatsStoreException(
+ "Cannot initialise SqliteStatsStore at " + _dbPath, secondEx);
+ }
+ }
+ }
+
+ private void openAndApplySchema()
+ throws Exception {
+ try {
+ openConnections();
+ } catch (Exception e) {
+ // Leave nothing bound to a database init() may be about to delete: a surviving connection
+ // would keep the old inode alive and serve reads from a file nothing writes to any more.
+ closeConnections();
+ throw e;
+ }
+ }
+
+ private void openConnections()
+ throws Exception {
+ Files.createDirectories(_dbDirectory);
+ String jdbcUrl = "jdbc:sqlite:" + _dbPath.toAbsolutePath();
+
+ // Open the shared writer connection.
+ // Set WAL and synchronous PRAGMAs with autoCommit=true (WAL mode change cannot be done
+ // inside a transaction), then switch to manual-commit mode for subsequent writes.
+ Connection conn = DriverManager.getConnection(jdbcUrl);
+ try (Statement st = conn.createStatement()) {
+ st.execute("PRAGMA journal_mode=WAL");
+ st.execute("PRAGMA synchronous=NORMAL");
+ }
+ applySchema(conn);
+ conn.setAutoCommit(false);
+ _writeConn = conn;
+
+ // Open read-only connections for the pool
+ for (int i = 0; i < _readPoolSize; i++) {
+ Connection rConn = DriverManager.getConnection(jdbcUrl);
+ rConn.setAutoCommit(true);
+ try (Statement st = rConn.createStatement()) {
+ st.execute("PRAGMA journal_mode=WAL");
+ st.execute("PRAGMA synchronous=NORMAL");
+ }
+ PooledReader reader = new PooledReader(rConn);
+ if (!_readPool.offer(reader)) {
+ reader.close();
+ }
+ }
+ }
+
+ /// Closes and forgets every connection this store holds. Safe to call repeatedly.
+ private void closeConnections() {
+ // A rebuild drops the file the cached rollups were computed from, so they cannot outlive it.
+ _aggregates.clear();
+ _versions.clear();
+ closeQuietly(_writeConn);
+ _writeConn = null;
+ PooledReader pooled;
+ while ((pooled = _readPool.poll()) != null) {
+ // Closes the cached statements as well as the connection.
+ pooled.close();
+ }
+ }
+
+ /// Creates the schema on an empty database, or verifies that an existing one matches
+ /// [#SCHEMA_VERSION].
+ ///
+ /// SQLite's own `user_version` pragma carries the version, so no bookkeeping table and no
+ /// migration library is needed for what is a derived, rebuildable cache. A mismatch throws, which
+ /// [#init()] turns into delete-and-recreate.
+ private static void applySchema(Connection conn)
+ throws SQLException {
+ try (Statement st = conn.createStatement()) {
+ int version;
+ try (ResultSet rs = st.executeQuery("PRAGMA user_version")) {
+ version = rs.next() ? rs.getInt(1) : 0;
+ }
+ if (version == SCHEMA_VERSION) {
+ return;
+ }
+ if (version != 0) {
+ throw new SQLException("Stats store schema version " + version + " does not match expected "
+ + SCHEMA_VERSION + "; the store will be rebuilt");
+ }
+ for (String ddl : SCHEMA_DDL) {
+ st.execute(ddl);
+ }
+ st.execute("PRAGMA user_version=" + SCHEMA_VERSION);
+ }
+ }
+
+ /// Deletes the SQLite DB file and its WAL / SHM siblings if they exist.
+ private void deleteDbFiles() {
+ tryDelete(_dbPath);
+ tryDelete(_dbDirectory.resolve(DB_FILE_NAME + "-wal"));
+ tryDelete(_dbDirectory.resolve(DB_FILE_NAME + "-shm"));
+ }
+
+ private static void tryDelete(Path p) {
+ try {
+ Files.deleteIfExists(p);
+ } catch (IOException e) {
+ LOGGER.warn("Could not delete {}: {}", p, e.getMessage());
+ }
+ }
+
+ @Override
+ public void close() {
+ _closed = true;
+ synchronized (_writeLock) {
+ closeQuietly(_writeConn);
+ _writeConn = null;
+ }
+ PooledReader r;
+ while ((r = _readPool.poll()) != null) {
+ r.close();
+ }
+ }
+
+ /// Returns the writer, or throws if the store closed between [#checkOpen()] and acquiring the
+ /// lock. Without this a writer that passed the flag check would NPE on a nulled connection, and
+ /// an NPE is neither rolled back here nor caught by the listeners that call this.
+ private Connection writeConnection()
+ throws StatsStoreException {
+ if (_writeConn == null) {
+ throw new StatsStoreException("SqliteStatsStore is closed");
+ }
+ return _writeConn;
+ }
+
+ private static void closeQuietly(@Nullable Connection conn) {
+ if (conn != null) {
+ try {
+ conn.close();
+ } catch (SQLException e) {
+ LOGGER.debug("Error closing connection", e);
+ }
+ }
+ }
+
+ // ---------------------------------------------------------------------------
+ // Write operations
+ // ---------------------------------------------------------------------------
+
+ @Override
+ public void upsertSegmentStats(String tableNameWithType, List<SegmentStatsRow> rows)
+ throws StatsStoreException {
+ checkOpen();
+ long now = System.currentTimeMillis();
+ synchronized (_writeLock) {
+ try {
+ try (PreparedStatement ps = writeConnection().prepareStatement(SQL_UPSERT_SEGMENT)) {
+ for (SegmentStatsRow row : rows) {
+ ps.setString(1, tableNameWithType);
+ ps.setString(2, row.segmentName());
+ ps.setLong(3, row.crc());
+ ps.setLong(4, row.totalDocs());
+ ps.setLong(5, row.sizeBytes());
+ ps.setLong(6, row.startTimeMs());
+ ps.setLong(7, row.endTimeMs());
+ ps.setInt(8, row.consuming() ? 1 : 0);
+ ps.setLong(9, now);
+ ps.addBatch();
+ }
+ ps.executeBatch();
+ }
+ writeConnection().commit();
+ invalidate(tableNameWithType);
+ } catch (SQLException e) {
+ rollbackQuietly(_writeConn);
+ invalidate(tableNameWithType);
+ throw new StatsStoreException("upsertSegmentStats failed for " + tableNameWithType, e);
+ }
+ }
+ }
+
+ @Override
+ public void upsertSegmentColumnStats(String tableNameWithType, List<SegmentColumnStatsRow> rows)
+ throws StatsStoreException {
+ checkOpen();
+ long now = System.currentTimeMillis();
+ synchronized (_writeLock) {
+ try {
+ try (PreparedStatement ps = writeConnection().prepareStatement(SQL_UPSERT_COL)) {
+ for (SegmentColumnStatsRow row : rows) {
+ ps.setString(1, tableNameWithType);
+ ps.setString(2, row.segmentName());
+ ps.setString(3, row.columnName());
+ ps.setLong(4, row.ndv());
+ ps.setString(5, row.minValue());
+ ps.setString(6, row.maxValue());
+ ps.setInt(7, row.minTrusted() ? 1 : 0);
+ ps.setDouble(8, row.avgBytesPerValue());
+ ps.setDouble(9, row.nullFraction());
+ ps.setString(10, row.valueType() == null ? null : row.valueType().name());
+ ps.setLong(11, now);
+ ps.addBatch();
+ }
+ ps.executeBatch();
+ }
+ writeConnection().commit();
+ } catch (SQLException e) {
+ rollbackQuietly(_writeConn);
+ throw new StatsStoreException(
+ "upsertSegmentColumnStats failed for " + tableNameWithType, e);
+ }
+ }
+ }
+
+ @Override
+ public void removeSegments(String tableNameWithType, Collection<String> segmentNames)
+ throws StatsStoreException {
+ checkOpen();
+ if (segmentNames.isEmpty()) {
+ return;
+ }
+ synchronized (_writeLock) {
+ try {
+ try (PreparedStatement psSeg = writeConnection().prepareStatement(SQL_DELETE_SEGMENT);
+ PreparedStatement psCol = writeConnection().prepareStatement(SQL_DELETE_COL)) {
+ for (String seg : segmentNames) {
+ psSeg.setString(1, tableNameWithType);
+ psSeg.setString(2, seg);
+ psSeg.addBatch();
+ psCol.setString(1, tableNameWithType);
+ psCol.setString(2, seg);
+ psCol.addBatch();
+ }
+ psSeg.executeBatch();
+ psCol.executeBatch();
+ }
+ writeConnection().commit();
+ invalidate(tableNameWithType);
+ } catch (SQLException e) {
+ rollbackQuietly(_writeConn);
+ invalidate(tableNameWithType);
+ throw new StatsStoreException("removeSegments failed for " + tableNameWithType, e);
+ }
+ }
+ }
+
+ @Override
+ public boolean hasConsumingSegments(String tableNameWithType)
+ throws StatsStoreException {
+ return aggregate(tableNameWithType)._hasConsuming;
+ }
+
+
+ @Override
+ public Set<String> getTables()
+ throws StatsStoreException {
+ checkOpen();
+ PooledReader reader = borrowReader();
+ try {
+ Set<String> tables = new HashSet<>();
+ PreparedStatement ps = reader.statement(SQL_LIST_TABLES);
+ try (ResultSet rs = ps.executeQuery()) {
+ while (rs.next()) {
+ tables.add(rs.getString(1));
+ }
+ }
+ return tables;
+ } catch (SQLException e) {
+ throw new StatsStoreException("getTables failed", e);
+ } finally {
+ returnReader(reader);
+ }
+ }
+
+ @Override
+ public void purgeTable(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ synchronized (_writeLock) {
+ try {
+ try (PreparedStatement psSeg = writeConnection().prepareStatement(SQL_PURGE_TABLE_SEG);
+ PreparedStatement psCol = writeConnection().prepareStatement(SQL_PURGE_TABLE_COL)) {
+ psSeg.setString(1, tableNameWithType);
+ psSeg.executeUpdate();
+ psCol.setString(1, tableNameWithType);
+ psCol.executeUpdate();
+ }
+ writeConnection().commit();
+ // Drop the cache entry outright rather than only stamping it: nothing is left to roll up.
+ _aggregates.remove(tableNameWithType);
+ _versions.remove(tableNameWithType);
+ } catch (SQLException e) {
+ rollbackQuietly(_writeConn);
+ invalidate(tableNameWithType);
+ throw new StatsStoreException("purgeTable failed for " + tableNameWithType, e);
+ }
+ }
+ }
+
+ @Override
+ public void purgeAll()
+ throws StatsStoreException {
+ checkOpen();
+ synchronized (_writeLock) {
+ try {
+ try (Statement st = writeConnection().createStatement()) {
+ st.execute(SQL_PURGE_ALL_SEG);
+ st.execute(SQL_PURGE_ALL_COL);
+ }
+ writeConnection().commit();
+ _aggregates.clear();
+ _versions.clear();
+ } catch (SQLException e) {
+ rollbackQuietly(_writeConn);
+ invalidateAll();
+ throw new StatsStoreException("purgeAll failed", e);
+ }
+ }
+ }
+
+ // ---------------------------------------------------------------------------
+ // Read operations
+ // ---------------------------------------------------------------------------
+
+ @Override
+ public Map<String, Long> getSegmentCrcs(String tableNameWithType)
+ throws StatsStoreException {
+ checkOpen();
+ PooledReader reader = borrowReader();
+ try {
+ Map<String, Long> result = new HashMap<>();
+ PreparedStatement ps = reader.statement(SQL_GET_CRCS);
+ {
+ ps.setString(1, tableNameWithType);
+ try (ResultSet rs = ps.executeQuery()) {
+ while (rs.next()) {
+ result.put(rs.getString(1), rs.getLong(2));
+ }
+ }
+ }
+ return result;
+ } catch (SQLException e) {
+ throw new StatsStoreException("getSegmentCrcs failed for " + tableNameWithType, e);
+ } finally {
+ returnReader(reader);
+ }
+ }
+
+ @Override
+ @Nullable
+ public TableStatistics getTableStats(String tableNameWithType)
+ throws StatsStoreException {
+ return aggregate(tableNameWithType)._stats;
+ }
+
+
+ /// Returns per-column statistics aggregated across all non-consuming segments for the given
+ /// table and column, or `null` if no rows exist.
+ ///
+ /// #### NDV
+ /// Returned as `MAX(ndv)` over segments with [StatConfidence#ESTIMATED]. The
+ /// true value lies in `[MAX(ndv), min(SUM(ndv), tableRowCount)]`; we report the lower
+ /// bound (MAX) because the upper bound is not representable as a single value in the contract.
+ ///
+ /// #### Min/Max
+ /// Compared numerically when both values parse as [Double], else lexically. Comparison
+ /// is done in Java (not SQL) to avoid SQLite TEXT-affinity ordering issues (e.g. "9" > "10").
+ @Override
+ @Nullable
+ public ColumnStatistics getColumnStats(String tableNameWithType, String columnName)
+ throws StatsStoreException {
+ checkOpen();
+ PooledReader reader = borrowReader();
+ try {
+ StatsAggregations.ColumnStatsAccumulator accumulator = new StatsAggregations.ColumnStatsAccumulator();
+ PreparedStatement ps = reader.statement(SQL_COL_STATS);
+ {
+ ps.setString(1, tableNameWithType);
+ ps.setString(2, columnName);
+ try (ResultSet rs = ps.executeQuery()) {
+ while (rs.next()) {
+ accumulator.add(rs.getLong(1), SegmentColumnStatsRow.builder()
+ .segmentName("")
+ .columnName(columnName)
+ .ndv(rs.getLong(2))
+ .bounds(rs.getString(3), rs.getString(4), ColumnValueType.fromName(rs.getString(8)))
+ .minTrusted(rs.getInt(5) != 0)
+ .avgBytesPerValue(rs.getDouble(6))
+ .nullFraction(rs.getDouble(7))
+ .build());
+ }
+ }
+ }
+ return accumulator.isEmpty() ? null : accumulator.build(columnName);
+ } catch (SQLException e) {
+ throw new StatsStoreException(
+ "getColumnStats failed for " + tableNameWithType + "." + columnName, e);
+ } finally {
+ returnReader(reader);
+ }
+ }
+
+ @Override
+ public OptionalLong estimateRowsInTimeRange(String tableNameWithType, long startMs, long endMs)
+ throws StatsStoreException {
+ checkOpen();
+ PooledReader reader = borrowReader();
+ try {
+ long totalRows = 0;
+ boolean hasAnyRow = false;
+
+ PreparedStatement ps = reader.statement(SQL_TIME_RANGE);
+ {
+ ps.setString(1, tableNameWithType);
+ ps.setLong(2, startMs);
+ ps.setLong(3, endMs);
+ ps.setString(4, tableNameWithType);
+ try (ResultSet rs = ps.executeQuery()) {
+ while (rs.next()) {
+ long docs = rs.getLong(1);
+ long segStart = rs.getLong(2);
+ long segEnd = rs.getLong(3);
+
+ hasAnyRow = true;
+
+ // Existence sentinel (see SQL_TIME_RANGE): committed stats exist, contributes 0 rows
+ if (segStart == -2) {
+ continue;
+ }
+
+ // The SQL predicate already prunes non-overlapping segments; the shared helper decides
+ // how much of a surviving segment counts (see StatsAggregations#overlapRows).
+ totalRows += StatsAggregations.overlapRows(docs, segStart, segEnd, startMs, endMs);
+ }
+ }
+ }
+
+ return hasAnyRow ? OptionalLong.of(totalRows) : OptionalLong.empty();
+ } catch (SQLException e) {
+ throw new StatsStoreException(
+ "estimateRowsInTimeRange failed for " + tableNameWithType, e);
+ } finally {
+ returnReader(reader);
+ }
+ }
+
+ // ---------------------------------------------------------------------------
+ // Helpers
+ // ---------------------------------------------------------------------------
+
+ private void checkOpen()
+ throws StatsStoreException {
+ if (_closed) {
+ throw new StatsStoreException("SqliteStatsStore is closed");
+ }
+ }
+
+ private PooledReader borrowReader()
+ throws StatsStoreException {
+ PooledReader reader = _readPool.poll();
+ if (reader != null) {
+ return reader;
+ }
+ // The pool is a bound, not a hint. Opening a connection per waiting reader would make the
+ // exhausted case -- the normal one on a broker with more planning threads than pooled
+ // connections -- the most expensive path: every read would pay a full database open plus its
+ // own page cache, and nothing would cap the file descriptors a burst can hold open. Waiting
+ // briefly instead keeps the cost bounded, and the caller already degrades to no statistics
+ // rather than failing a query.
+ try {
+ reader = _readPool.poll(_readBorrowTimeoutMs, TimeUnit.MILLISECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new StatsStoreException("Interrupted while waiting for a read connection", e);
+ }
+ if (reader == null) {
+ throw new StatsStoreException(
+ "Timed out after " + _readBorrowTimeoutMs + "ms waiting for one of " + _readPoolSize
+ + " read connections");
+ }
+ // A closed store must not hand out a connection it is in the middle of tearing down.
+ checkOpen();
+ return reader;
+ }
+
+ private void returnReader(@Nullable PooledReader reader) {
+ if (reader == null) {
+ return;
+ }
+ // Do not re-pool into a store that closed while this read was in flight, or close() would
+ // leave a live connection -- and its statements -- behind.
+ if (_closed || !_readPool.offer(reader)) {
+ reader.close();
+ }
+ }
+
+ /// A pooled read connection together with the statements compiled on it.
+ ///
+ /// SQLite compiles each statement into a VDBE program, and that compilation runs the query
+ /// planner -- so re-preparing on every call re-does index selection. Measured on this schema:
+ /// preparing is ~58% of total query cost for a table with 50 segments, where compilation
+ /// dominates execution, falling to ~4% at 1000 segments. Most tables sit at the small end and
+ /// the planner issues several of these per compile, so the statements are kept with their
+ /// connection rather than rebuilt.
+ ///
+ /// Not thread-safe, and does not need to be: an instance is owned by exactly one thread between
+ /// [#borrowReader()] and [#returnReader].
+ private static final class PooledReader {
+ private final Connection _conn;
+ private final Map<String, PreparedStatement> _statements = new HashMap<>();
+
+ PooledReader(Connection conn) {
+ _conn = conn;
+ }
+
+ /// The statement for `sql`, compiled on first use and reused afterwards, with any parameters
+ /// from the previous execution cleared.
+ PreparedStatement statement(String sql)
+ throws SQLException {
+ PreparedStatement ps = _statements.get(sql);
+ if (ps == null) {
+ ps = _conn.prepareStatement(sql);
+ _statements.put(sql, ps);
+ }
+ ps.clearParameters();
+ return ps;
+ }
+
+ void close() {
+ for (PreparedStatement ps : _statements.values()) {
+ try {
+ ps.close();
+ } catch (SQLException e) {
+ LOGGER.debug("Closing pooled statement failed", e);
+ }
+ }
+ _statements.clear();
+ closeQuietly(_conn);
+ }
+ }
+
+ private static void rollbackQuietly(@Nullable Connection conn) {
+ if (conn == null) {
+ return;
+ }
+ try {
+ conn.rollback();
+ } catch (SQLException e) {
+ LOGGER.debug("Rollback failed", e);
+ }
+ }
+
+ /// Returns the cached rollup for a table, recomputing it only when a write has landed since the
+ /// cached copy was taken.
+ ///
+ /// Query planning asks for these on every compile while writes arrive at segment-push cadence,
+ /// so without a cache every plan would scan all of a table's segment rows -- on a table with
+ /// hundreds of thousands of segments that lands directly in planning latency. The version stamp
+ /// is what makes caching safe without holding the write lock: a rollup computed from rows that
+ /// changed underneath it is used for this call but not published.
+ private CachedAggregate aggregate(String tableNameWithType)
+ throws StatsStoreException {
+ AtomicLong version = _versions.computeIfAbsent(tableNameWithType, k -> new AtomicLong());
+ long observed = version.get();
+ CachedAggregate cached = _aggregates.get(tableNameWithType);
+ if (cached != null && cached._version == observed) {
+ return cached;
+ }
+ CachedAggregate computed = computeAggregate(tableNameWithType, observed);
+ if (version.get() == observed) {
+ _aggregates.put(tableNameWithType, computed);
+ }
+ return computed;
+ }
+
+ private CachedAggregate computeAggregate(String tableNameWithType, long version)
+ throws StatsStoreException {
+ checkOpen();
+ PooledReader reader = borrowReader();
+ try {
+ PreparedStatement ps = reader.statement(SQL_TABLE_STATS);
+ ps.setString(1, tableNameWithType);
+ try (ResultSet rs = ps.executeQuery()) {
+ // Aggregates always return one row, holding SQL NULLs when the table has none.
+ if (!rs.next()) {
+ return new CachedAggregate(version, null, false);
+ }
+ long totalDocs = rs.getLong(1);
+ long sizeBytes = rs.getLong(2);
+ long maxUpdatedAt = rs.getLong(3);
+ long committed = rs.getLong(4);
+ boolean hasConsuming = rs.getLong(5) == 1;
+ if (committed == 0) {
+ return new CachedAggregate(version, null, hasConsuming);
+ }
+ return new CachedAggregate(version, TableStatistics.builder()
+ .rowCount(totalDocs, StatConfidence.EXACT)
+ .tableSizeBytes(sizeBytes, StatConfidence.EXACT)
+ .updatedAtMs(maxUpdatedAt)
+ .build(), hasConsuming);
+ }
+ } catch (SQLException e) {
+ throw new StatsStoreException("getTableStats failed for " + tableNameWithType, e);
+ } finally {
+ returnReader(reader);
+ }
+ }
+
+ /// Marks a table's cached rollup stale. Called from the write paths, which hold `_writeLock`.
+ private void invalidate(String tableNameWithType) {
+ _versions.computeIfAbsent(tableNameWithType, k -> new AtomicLong()).incrementAndGet();
+ _aggregates.remove(tableNameWithType);
+ }
+
+ /// Marks every table's cached rollup stale, for writes that span tables.
+ private void invalidateAll() {
+ _aggregates.clear();
+ _versions.values().forEach(AtomicLong::incrementAndGet);
+ }
+
+ private static final class CachedAggregate {
+ private final long _version;
+ @Nullable
+ private final TableStatistics _stats;
+ private final boolean _hasConsuming;
+
+ CachedAggregate(long version, @Nullable TableStatistics stats, boolean hasConsuming) {
+ _version = version;
+ _stats = stats;
+ _hasConsuming = hasConsuming;
+ }
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/stats/InMemoryStatsStoreTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/InMemoryStatsStoreTest.java
new file mode 100644
index 00000000000..c55086c390b
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/InMemoryStatsStoreTest.java
@@ -0,0 +1,52 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.stats;
+
+import java.util.List;
+import org.apache.pinot.query.planner.spi.stats.SegmentStatsRow;
+import org.apache.pinot.query.planner.spi.stats.StatsStore;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertTrue;
+
+
+/// Runs the shared [StatsStore] contract against [InMemoryStatsStore], plus the one behavior that
+/// distinguishes it: a new instance starts empty, so the caller re-collects everything.
+public class InMemoryStatsStoreTest extends StatsStoreContractTest {
+
+ @Override
+ protected StatsStore createStore() {
+ return new InMemoryStatsStore();
+ }
+
+ /// A fresh store keeps nothing from a previous one. Reconciliation relies on this: an empty
+ /// crc map makes every segment look new, so the listener re-upserts the lot.
+ @Test
+ public void testNothingSurvivesANewInstance()
+ throws Exception {
+ _store.upsertSegmentStats("myTable_OFFLINE",
+ List.of(new SegmentStatsRow("seg1", 42L, 500L, 2000L, 0L, 100L, false)));
+ assertTrue(_store.getSegmentCrcs("myTable_OFFLINE").containsKey("seg1"));
+
+ try (StatsStore fresh = createStore()) {
+ fresh.init();
+ assertTrue(fresh.getSegmentCrcs("myTable_OFFLINE").isEmpty(), "A new store must start empty");
+ }
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/stats/SqliteStatsStoreTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/SqliteStatsStoreTest.java
new file mode 100644
index 00000000000..e5362b488cf
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/SqliteStatsStoreTest.java
@@ -0,0 +1,259 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.stats;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.Statement;
+import java.util.ArrayDeque;
+import java.util.Deque;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.BlockingQueue;
+import org.apache.pinot.query.planner.spi.stats.SegmentStatsRow;
+import org.apache.pinot.query.planner.spi.stats.StatsStore;
+import org.apache.pinot.query.planner.spi.stats.StatsStoreException;
+import org.apache.pinot.query.planner.spi.stats.TableStatistics;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.*;
+
+
+/// Runs the shared [StatsStore] contract against [SqliteStatsStore], plus the cases that only a
+/// durable, file-backed store can have.
+public class SqliteStatsStoreTest extends StatsStoreContractTest {
+
+ private Path _tempDir;
+
+ @Override
+ protected StatsStore createStore()
+ throws Exception {
+ _tempDir = Files.createTempDirectory("sqlite-stats-test-");
+ return new SqliteStatsStore(_tempDir);
+ }
+
+ @Override
+ protected void cleanUp()
+ throws Exception {
+ deleteRecursively(_tempDir);
+ }
+
+ // ---------------------------------------------------------------------------
+ // Persistence across reopen
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testPersistenceAcrossReopen()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("seg1", 42L, 500L, 2000L, 0L, 100L, false)
+ ));
+ _store.close();
+ _store = null;
+
+ // Open a new store on the same directory
+ SqliteStatsStore store2 = new SqliteStatsStore(_tempDir);
+ store2.init();
+ try {
+ Map<String, Long> crcs = store2.getSegmentCrcs(TABLE_A);
+ assertEquals(crcs.size(), 1);
+ assertEquals(crcs.get("seg1").longValue(), 42L);
+
+ TableStatistics stats = store2.getTableStats(TABLE_A);
+ assertNotNull(stats);
+ assertEquals(stats.getRowCount(), 500L);
+ } finally {
+ store2.close();
+ }
+ }
+
+ // ---------------------------------------------------------------------------
+ // Corruption recovery
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testCorruptionRecovery()
+ throws Exception {
+ // Write some data, then close
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("seg1", 1L, 100L, 1000L, 0L, 10L, false)
+ ));
+ _store.close();
+ _store = null;
+
+ // Corrupt the DB file
+ Path dbFile = _tempDir.resolve("broker-stats.sqlite");
+ Files.write(dbFile, "this is not a valid sqlite database file!!!".getBytes());
+
+ // Opening a new store should recover silently
+ SqliteStatsStore store2 = new SqliteStatsStore(_tempDir);
+ store2.init(); // must not throw
+ try {
+ // After recovery, the store is empty
+ assertNull(store2.getTableStats(TABLE_A));
+ Map<String, Long> crcs = store2.getSegmentCrcs(TABLE_A);
+ assertTrue(crcs.isEmpty());
+ } finally {
+ store2.close();
+ }
+ }
+
+ /// The schema carries no migration path on purpose: rows here are derived from ZooKeeper and can
+ /// always be re-collected, so a store written by a different schema version is discarded rather
+ /// than migrated -- the same recovery path a corrupt file takes.
+ @Test
+ public void testSchemaVersionMismatchRebuilds()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("seg1", 42L, 500L, 2000L, 0L, 100L, false)));
+ _store.close();
+ _store = null;
+
+ Path dbFile = _tempDir.resolve("broker-stats.sqlite");
+ try (Connection conn = DriverManager.getConnection("jdbc:sqlite:" + dbFile.toAbsolutePath());
+ Statement st = conn.createStatement()) {
+ st.execute("PRAGMA user_version=9999");
+ }
+
+ try (SqliteStatsStore reopened = new SqliteStatsStore(_tempDir)) {
+ reopened.init();
+ assertTrue(reopened.getSegmentCrcs(TABLE_A).isEmpty(), "A foreign schema version must be rebuilt empty");
+ assertNull(reopened.getTableStats(TABLE_A));
+ }
+ }
+
+ // ---------------------------------------------------------------------------
+ // Read pool sizing and borrow timeout
+ // ---------------------------------------------------------------------------
+
+ /// The pool is a concurrency bound, not a per-caller connection, so a store configured with a
+ /// single reader must still serve every caller correctly: readers are borrowed and returned,
+ /// never leaked and never handed to two callers at once.
+ @Test
+ public void testSingleReaderPoolStillServesEveryRead()
+ throws Exception {
+ Path dir = Files.createTempDirectory("sqlite-stats-pool-");
+ try (SqliteStatsStore store = new SqliteStatsStore(dir, 1, 5000)) {
+ store.init();
+ store.upsertSegmentStats(TABLE_A, List.of(seg("seg1", 1L, 500L, 2000L, 0L, 100L, false)));
+
+ // More reads than pooled readers: each must borrow, use and return the one connection.
+ for (int i = 0; i < 8; i++) {
+ TableStatistics stats = store.getTableStats(TABLE_A);
+ assertNotNull(stats, "Read " + i + " got no statistics");
+ assertEquals(stats.getRowCount(), 500L, "Read " + i + " saw the wrong row count");
+ }
+ } finally {
+ deleteRecursively(dir);
+ }
+ }
+
+ /// An exhausted pool must give up after the configured timeout rather than block a planner
+ /// thread indefinitely, and must recover once readers come back.
+ @Test
+ public void testBorrowTimesOutWhenPoolIsExhausted()
+ throws Exception {
+ Path dir = Files.createTempDirectory("sqlite-stats-timeout-");
+ long timeoutMs = 100;
+ try (SqliteStatsStore store = new SqliteStatsStore(dir, 2, timeoutMs)) {
+ store.init();
+ store.upsertSegmentStats(TABLE_A, List.of(seg("seg1", 1L, 500L, 2000L, 0L, 100L, false)));
+
+ // Hold every reader. Reflection rather than blocked threads: an empty pool is exactly the
+ // state under test, and reproducing it with concurrent reads would race against how fast
+ // SQLite answers.
+ Deque<Object> held = drainReadPool(store);
+ assertEquals(held.size(), 2, "Expected the configured number of pooled readers");
+ try {
+ long startNanos = System.nanoTime();
+ StatsStoreException e = expectThrows(StatsStoreException.class, () -> store.getTableStats(TABLE_A));
+ long elapsedMs = (System.nanoTime() - startNanos) / 1_000_000;
+
+ assertTrue(e.getMessage().contains("Timed out"), "Unexpected message: " + e.getMessage());
+ assertTrue(elapsedMs >= timeoutMs,
+ "Gave up after " + elapsedMs + "ms, before the configured " + timeoutMs + "ms");
+ } finally {
+ returnToReadPool(store, held);
+ }
+
+ // With the readers back, reads succeed again.
+ TableStatistics recovered = store.getTableStats(TABLE_A);
+ assertNotNull(recovered);
+ assertEquals(recovered.getRowCount(), 500L);
+ } finally {
+ deleteRecursively(dir);
+ }
+ }
+
+ private static Deque<Object> drainReadPool(SqliteStatsStore store)
+ throws Exception {
+ BlockingQueue<Object> pool = readPool(store);
+ Deque<Object> held = new ArrayDeque<>();
+ Object reader;
+ while ((reader = pool.poll()) != null) {
+ held.add(reader);
+ }
+ return held;
+ }
+
+ private static void returnToReadPool(SqliteStatsStore store, Deque<Object> held)
+ throws Exception {
+ BlockingQueue<Object> pool = readPool(store);
+ for (Object reader : held) {
+ pool.offer(reader);
+ }
+ }
+
+ @SuppressWarnings("unchecked")
+ private static BlockingQueue<Object> readPool(SqliteStatsStore store)
+ throws Exception {
+ java.lang.reflect.Field field = SqliteStatsStore.class.getDeclaredField("_readPool");
+ field.setAccessible(true);
+ return (BlockingQueue<Object>) field.get(store);
+ }
+
+ // ---------------------------------------------------------------------------
+ // Factory helpers
+ // ---------------------------------------------------------------------------
+
+ private static SegmentStatsRow seg(String name, long crc, long docs, long sizeBytes,
+ long startMs, long endMs, boolean consuming) {
+ return new SegmentStatsRow(name, crc, docs, sizeBytes, startMs, endMs, consuming);
+ }
+
+ /// Recursively deletes a directory tree.
+ private static void deleteRecursively(Path dir)
+ throws IOException {
+ if (dir == null || !Files.exists(dir)) {
+ return;
+ }
+ try (var stream = Files.walk(dir)) {
+ stream.sorted(java.util.Comparator.reverseOrder())
+ .forEach(p -> {
+ try {
+ Files.deleteIfExists(p);
+ } catch (IOException e) {
+ // ignore
+ }
+ });
+ }
+ }
+}
diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/stats/StatsStoreContractTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/StatsStoreContractTest.java
new file mode 100644
index 00000000000..08bcbdfda4c
--- /dev/null
+++ b/pinot-broker/src/test/java/org/apache/pinot/broker/stats/StatsStoreContractTest.java
@@ -0,0 +1,560 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.broker.stats;
+
+import java.util.List;
+import java.util.Map;
+import java.util.OptionalLong;
+import java.util.Set;
+import org.apache.pinot.query.planner.spi.stats.ColumnStatistics;
+import org.apache.pinot.query.planner.spi.stats.ColumnValueType;
+import org.apache.pinot.query.planner.spi.stats.SegmentColumnStatsRow;
+import org.apache.pinot.query.planner.spi.stats.SegmentStatsRow;
+import org.apache.pinot.query.planner.spi.stats.StatConfidence;
+import org.apache.pinot.query.planner.spi.stats.StatsAggregations;
+import org.apache.pinot.query.planner.spi.stats.StatsStore;
+import org.apache.pinot.query.planner.spi.stats.StatsStoreException;
+import org.apache.pinot.query.planner.spi.stats.TableStatistics;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.*;
+
+
+/// Behavior every [StatsStore] implementation must share.
+///
+/// The optimizer must not change with the configured store, so these cases are defined once here
+/// and run against each implementation by a subclass. Anything specific to one implementation
+/// (durability, corruption recovery) belongs in that subclass instead.
+public abstract class StatsStoreContractTest {
+ protected static final String TABLE_A = "myTable_OFFLINE";
+ protected static final String TABLE_B = "otherTable_REALTIME";
+
+ protected StatsStore _store;
+
+ /// Creates a fresh, un-initialized store for one test method.
+ protected abstract StatsStore createStore()
+ throws Exception;
+
+ /// Releases anything the concrete store needed beyond [StatsStore#close()].
+ protected void cleanUp()
+ throws Exception {
+ }
+
+ @BeforeMethod
+ public void setUp()
+ throws Exception {
+ _store = createStore();
+ _store.init();
+ }
+
+ @AfterMethod
+ public void tearDown()
+ throws Exception {
+ if (_store != null) {
+ _store.close();
+ }
+ cleanUp();
+ }
+
+ // ---------------------------------------------------------------------------
+ // Round-trip: upsert + getTableStats
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testRoundTripTableStats()
+ throws Exception {
+ List<SegmentStatsRow> rowsA = List.of(
+ seg("seg1", 100L, 1000L, 5000L, 0L, 100L, false),
+ seg("seg2", 200L, 2000L, 8000L, 100L, 200L, false)
+ );
+ List<SegmentStatsRow> rowsB = List.of(
+ seg("seg3", 300L, 500L, 1000L, 0L, 50L, false)
+ );
+
+ _store.upsertSegmentStats(TABLE_A, rowsA);
+ _store.upsertSegmentStats(TABLE_B, rowsB);
+
+ TableStatistics statsA = _store.getTableStats(TABLE_A);
+ assertNotNull(statsA);
+ assertEquals(statsA.getRowCount(), 3000L);
+ assertEquals(statsA.getTableSizeBytes(), 13000L);
+ assertEquals(statsA.getRowCountConfidence(), StatConfidence.EXACT);
+ assertEquals(statsA.getSizeConfidence(), StatConfidence.EXACT);
+
+ TableStatistics statsB = _store.getTableStats(TABLE_B);
+ assertNotNull(statsB);
+ assertEquals(statsB.getRowCount(), 500L);
+ }
+
+ @Test
+ public void testConsumingSegmentsExcludedFromTableStats()
+ throws Exception {
+ List<SegmentStatsRow> rows = List.of(
+ seg("committed", 111L, 1000L, 5000L, 0L, 100L, false),
+ seg("consuming", 222L, 999L, 1000L, 100L, 200L, true)
+ );
+ _store.upsertSegmentStats(TABLE_A, rows);
+
+ TableStatistics stats = _store.getTableStats(TABLE_A);
+ assertNotNull(stats);
+ // Only the committed segment should be counted
+ assertEquals(stats.getRowCount(), 1000L);
+ assertEquals(stats.getTableSizeBytes(), 5000L);
+ }
+
+ @Test
+ public void testGetTableStatsNullWhenEmpty()
+ throws Exception {
+ assertNull(_store.getTableStats(TABLE_A));
+ }
+
+ // ---------------------------------------------------------------------------
+ // Upsert overwrite
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testUpsertOverwrite()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A,
+ List.of(seg("seg1", 100L, 1000L, 5000L, 0L, 100L, false)));
+
+ // Overwrite with new values
+ _store.upsertSegmentStats(TABLE_A,
+ List.of(seg("seg1", 999L, 2500L, 9000L, 0L, 100L, false)));
+
+ Map<String, Long> crcs = _store.getSegmentCrcs(TABLE_A);
+ assertEquals(crcs.size(), 1);
+ assertEquals(crcs.get("seg1").longValue(), 999L);
+
+ TableStatistics stats = _store.getTableStats(TABLE_A);
+ assertNotNull(stats);
+ assertEquals(stats.getRowCount(), 2500L);
+ }
+
+ // ---------------------------------------------------------------------------
+ // removeSegments
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testRemoveSegments()
+ throws Exception {
+ List<SegmentStatsRow> rows = List.of(
+ seg("seg1", 1L, 100L, 1000L, 0L, 10L, false),
+ seg("seg2", 2L, 200L, 2000L, 10L, 20L, false)
+ );
+ _store.upsertSegmentStats(TABLE_A, rows);
+ _store.upsertSegmentColumnStats(TABLE_A, List.of(
+ col("seg1", "colA", 10L, "1", "9", true, 4.0, 0.0),
+ col("seg2", "colA", 20L, "2", "8", true, 4.0, 0.0)
+ ));
+
+ _store.removeSegments(TABLE_A, List.of("seg1"));
+
+ Map<String, Long> crcs = _store.getSegmentCrcs(TABLE_A);
+ assertFalse(crcs.containsKey("seg1"));
+ assertTrue(crcs.containsKey("seg2"));
+
+ // Column stats for seg1 should also be removed
+ ColumnStatistics colStats = _store.getColumnStats(TABLE_A, "colA");
+ assertNotNull(colStats);
+ // Only seg2's stats remain: ndv=20, min=2, max=8
+ assertEquals(colStats.getNdv(), 20L);
+ }
+
+ // ---------------------------------------------------------------------------
+ // Column stats aggregation
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testColumnStatsAggregation()
+ throws Exception {
+ // 3 segments with different stats; docs used for weighting
+ List<SegmentStatsRow> segs = List.of(
+ seg("s1", 1L, 100L, 1000L, 0L, 100L, false),
+ seg("s2", 2L, 200L, 2000L, 100L, 200L, false),
+ seg("s3", 3L, 300L, 3000L, 200L, 300L, false)
+ );
+ _store.upsertSegmentStats(TABLE_A, segs);
+
+ // ndv: 10, 20, 30 → MAX = 30
+ // min: "9", "10", "2" → numeric min = 2 (i.e. "2")
+ // max: "9", "10", "20" → numeric max = 20 (i.e. "20")
+ // minTrusted: true, false, true → AND = false
+ // avgBytes: 4.0, 8.0, 2.0 (weighted by docs 100/200/300)
+ // weighted = (4*100 + 8*200 + 2*300) / 600 = (400+1600+600)/600 = 2600/600 ≈ 4.333...
+ // nullFraction: 0.1, 0.2, 0.0 (weighted)
+ // weighted = (0.1*100 + 0.2*200 + 0.0*300) / 600 = (10+40+0)/600 = 50/600 ≈ 0.0833...
+ List<SegmentColumnStatsRow> cols = List.of(
+ col("s1", "colA", 10L, "9", "9", true, 4.0, 0.1),
+ col("s2", "colA", 20L, "10", "10", false, 8.0, 0.2),
+ col("s3", "colA", 30L, "2", "20", true, 2.0, 0.0)
+ );
+ _store.upsertSegmentColumnStats(TABLE_A, cols);
+
+ ColumnStatistics cs = _store.getColumnStats(TABLE_A, "colA");
+ assertNotNull(cs);
+
+ assertEquals(cs.getNdv(), 30L);
+ assertEquals(cs.getNdvConfidence(), StatConfidence.ESTIMATED);
+
+ // Ordered as LONG (the recorded type), and returned as the Java type that ordering implies:
+ // min of "9","10","2" is 2, max of "9","10","20" is 20.
+ assertEquals(cs.getMinValue(), 2L);
+ assertEquals(cs.getMaxValue(), 20L);
+
+ // minTrusted: s2 is false → overall false
+ assertFalse(cs.isMinTrusted());
+
+ // Weighted avgBytes ≈ 4.333
+ assertEquals(cs.getAvgBytesPerValue(), 2600.0 / 600.0, 0.001);
+
+ // Weighted nullFraction ≈ 0.0833
+ assertEquals(cs.getNullFraction(), 50.0 / 600.0, 0.001);
+ }
+
+ @Test
+ public void testNumericMinMaxOrdering()
+ throws Exception {
+ // "9" vs "10": lexically "9" > "10", numerically "9" < "10"
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 100L, 1000L, 0L, 100L, false),
+ seg("s2", 2L, 100L, 1000L, 0L, 100L, false)
+ ));
+ _store.upsertSegmentColumnStats(TABLE_A, List.of(
+ col("s1", "colA", 5L, "9", "9", true, 4.0, 0.0),
+ col("s2", "colA", 5L, "10", "10", true, 4.0, 0.0)
+ ));
+
+ ColumnStatistics cs = _store.getColumnStats(TABLE_A, "colA");
+ assertNotNull(cs);
+ // "9" vs "10" orders numerically under LONG, where text order would put "10" first.
+ assertEquals(cs.getMinValue(), 9L);
+ assertEquals(cs.getMaxValue(), 10L);
+ }
+
+ @Test
+ public void testColumnStatsNullWhenNoRows()
+ throws Exception {
+ assertNull(_store.getColumnStats(TABLE_A, "colA"));
+ }
+
+ // ---------------------------------------------------------------------------
+ // estimateRowsInTimeRange
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testEstimateRowsFullOverlap()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, 0L, 100L, false), // [0, 100)
+ seg("s2", 2L, 2000L, 0L, 100L, 200L, false) // [100, 200)
+ ));
+
+ // Query [0, 200) → both fully inside
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 0L, 200L);
+ assertTrue(result.isPresent());
+ assertEquals(result.getAsLong(), 3000L);
+ }
+
+ @Test
+ public void testEstimateRowsPartialOverlap()
+ throws Exception {
+ // Segment [0, 100), docs=1000
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, 0L, 100L, false)
+ ));
+
+ // Query [25, 75) → 50% overlap → 500 rows
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 25L, 75L);
+ assertTrue(result.isPresent());
+ assertEquals(result.getAsLong(), 500L);
+ }
+
+ @Test
+ public void testEstimateRowsNoOverlap()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, 0L, 100L, false)
+ ));
+
+ // Query [200, 300) → no overlap
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 200L, 300L);
+ assertTrue(result.isPresent());
+ assertEquals(result.getAsLong(), 0L);
+ }
+
+ @Test
+ public void testEstimateRowsUnknownTimesSegment()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, -1L, -1L, false)
+ ));
+
+ // Unknown times → conservative, always include
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 0L, 100L);
+ assertTrue(result.isPresent());
+ assertEquals(result.getAsLong(), 1000L);
+ }
+
+ @Test
+ public void testEstimateRowsEmptyOptionalWhenNoRows()
+ throws Exception {
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 0L, 100L);
+ assertFalse(result.isPresent());
+ }
+
+ @Test
+ public void testEstimateRowsEmptyOptionalWhenOnlyConsumingSegments()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, 0L, 100L, true)
+ ));
+
+ // Consuming segments are excluded from time-range estimates; with no committed segments the
+ // result must be empty ("no stats"), never of(0) ("provably zero rows in range").
+ OptionalLong result = _store.estimateRowsInTimeRange(TABLE_A, 0L, 100L);
+ assertFalse(result.isPresent());
+ }
+
+ @Test
+ public void testEstimateRowsBoundaryAdjacency()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("s1", 1L, 1000L, 0L, 100L, 200L, false) // [100, 200)
+ ));
+
+ // Segment ending exactly at the range start → no overlap (half-open semantics)
+ OptionalLong before = _store.estimateRowsInTimeRange(TABLE_A, 200L, 300L);
+ assertTrue(before.isPresent());
+ assertEquals(before.getAsLong(), 0L);
+
+ // Segment starting exactly at the range end → no overlap
+ OptionalLong after = _store.estimateRowsInTimeRange(TABLE_A, 0L, 100L);
+ assertTrue(after.isPresent());
+ assertEquals(after.getAsLong(), 0L);
+
+ // Touching on both sides at once: range exactly equal to the segment → full overlap
+ OptionalLong exact = _store.estimateRowsInTimeRange(TABLE_A, 100L, 200L);
+ assertTrue(exact.isPresent());
+ assertEquals(exact.getAsLong(), 1000L);
+ }
+
+ // ---------------------------------------------------------------------------
+ // purgeTable / purgeAll
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testPurgeTable()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("seg1", 1L, 100L, 1000L, 0L, 10L, false)
+ ));
+ _store.upsertSegmentStats(TABLE_B, List.of(
+ seg("seg2", 2L, 200L, 2000L, 0L, 10L, false)
+ ));
+ _store.upsertSegmentColumnStats(TABLE_A, List.of(
+ col("seg1", "colA", 5L, "1", "9", true, 4.0, 0.0)
+ ));
+
+ _store.purgeTable(TABLE_A);
+
+ assertNull(_store.getTableStats(TABLE_A));
+ assertNull(_store.getColumnStats(TABLE_A, "colA"));
+ // TABLE_B should still be there
+ assertNotNull(_store.getTableStats(TABLE_B));
+ }
+
+ @Test
+ public void testPurgeAll()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(
+ seg("seg1", 1L, 100L, 1000L, 0L, 10L, false)
+ ));
+ _store.upsertSegmentStats(TABLE_B, List.of(
+ seg("seg2", 2L, 200L, 2000L, 0L, 10L, false)
+ ));
+
+ _store.purgeAll();
+
+ assertNull(_store.getTableStats(TABLE_A));
+ assertNull(_store.getTableStats(TABLE_B));
+ }
+
+ // ---------------------------------------------------------------------------
+ // Consuming segments
+ // ---------------------------------------------------------------------------
+
+ /// Drives whether realtime row counts are trusted (see `LogicalTableStatsResolver`), so both
+ /// stores must agree on it.
+ @Test
+ public void testHasConsumingSegments()
+ throws Exception {
+ assertFalse(_store.hasConsumingSegments(TABLE_A), "Unknown table has no consuming segments");
+
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("committed", 1L, 100L, 1000L, 0L, 10L, false)));
+ assertFalse(_store.hasConsumingSegments(TABLE_A));
+
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("consuming", 2L, 50L, 500L, 10L, 20L, true)));
+ assertTrue(_store.hasConsumingSegments(TABLE_A));
+
+ _store.removeSegments(TABLE_A, List.of("consuming"));
+ assertFalse(_store.hasConsumingSegments(TABLE_A), "Removing the consuming segment clears the flag");
+ }
+
+ /// A table with nothing but consuming segments has no trustworthy aggregate: reads must report
+ /// "no statistics" rather than a zero, and the segment must still be reconcilable by crc.
+ @Test
+ public void testOnlyConsumingSegmentsYieldNoTableOrColumnStats()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("consuming", 7L, 999L, 500L, 0L, 10L, true)));
+ _store.upsertSegmentColumnStats(TABLE_A,
+ List.of(col("consuming", "colA", 5L, "1", "9", true, 4.0, 0.0)));
+
+ assertNull(_store.getTableStats(TABLE_A), "Consuming-only table has no committed rows to aggregate");
+ assertNull(_store.getColumnStats(TABLE_A, "colA"), "A column row on a consuming segment does not count");
+ assertTrue(_store.getSegmentCrcs(TABLE_A).containsKey("consuming"), "crcs still cover consuming segments");
+ }
+
+ /// A column row whose segment row is gone contributes nothing: it supplies the doc count used for
+ /// weighting, so counting it would weight by a stale or absent value.
+ @Test
+ public void testColumnStatsIgnoreRowsWithoutASegment()
+ throws Exception {
+ _store.upsertSegmentColumnStats(TABLE_A,
+ List.of(col("ghost", "colA", 5L, "1", "9", true, 4.0, 0.0)));
+
+ assertNull(_store.getColumnStats(TABLE_A, "colA"));
+ }
+
+ // ---------------------------------------------------------------------------
+ // Shared aggregation semantics
+ // ---------------------------------------------------------------------------
+
+ /// Pins the ordering rules both stores rely on. The type has to be recorded with the values:
+ /// guessing "numeric if it parses" orders a STRING column numerically and rounds a LONG past
+ /// 2^53 — inwards, which would exclude rows that exist.
+ @Test
+ public void testValueOrderingIsDrivenByTheRecordedType() {
+ assertTrue(ColumnValueType.LONG.compare("9", "10") < 0, "LONG orders numerically");
+ assertTrue(ColumnValueType.STRING.compare("9", "10") > 0, "STRING orders lexically");
+
+ // 2^53 + 1 and 2^53 + 3 are indistinguishable as doubles; as LONG they are not.
+ assertTrue(ColumnValueType.LONG.compare("9007199254740993", "9007199254740995") < 0,
+ "LONG must not lose precision past 2^53");
+ assertEquals(ColumnValueType.LONG.toComparable("9007199254740993"), 9007199254740993L);
+ assertEquals(ColumnValueType.BIG_DECIMAL.toComparable("1.5"), new java.math.BigDecimal("1.5"));
+ assertEquals(ColumnValueType.STRING.toComparable("abc"), "abc");
+ assertNull(ColumnValueType.STRING.toComparable(null));
+
+ assertEquals(StatsAggregations.minOf(null, "5", ColumnValueType.LONG), "5", "null means unknown");
+ assertEquals(StatsAggregations.maxOf("5", null, ColumnValueType.LONG), "5");
+ assertNull(StatsAggregations.minOf(null, null, ColumnValueType.LONG));
+ }
+
+ /// Without a recorded type there is no defined ordering, so the bounds must be reported
+ /// untrusted rather than ordered on a guess — a consumer may prune on trusted bounds.
+ @Test
+ public void testUnknownValueTypeYieldsUntrustedBounds()
+ throws Exception {
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("s1", 1L, 100L, 1000L, 0L, 10L, false)));
+ _store.upsertSegmentColumnStats(TABLE_A,
+ List.of(new SegmentColumnStatsRow("s1", "colA", 5L, "9", "10", true, 4.0, 0.0, null)));
+
+ ColumnStatistics stats = _store.getColumnStats(TABLE_A, "colA");
+ assertNotNull(stats);
+ assertFalse(stats.isMinTrusted(), "Bounds with no recorded ordering must not be trusted");
+ }
+
+ /// Pins the per-segment overlap rules, including the cases a query never makes obvious.
+ @Test
+ public void testOverlapSemantics() {
+ assertEquals(StatsAggregations.overlapRows(100L, -1L, -1L, 0L, 10L), 100L, "Unknown times count in full");
+ assertEquals(StatsAggregations.overlapRows(100L, 10L, 20L, 0L, 100L), 100L, "Full containment");
+ assertEquals(StatsAggregations.overlapRows(100L, 100L, 200L, 0L, 100L), 0L, "Half-open upper bound excludes");
+ assertEquals(StatsAggregations.overlapRows(1000L, 0L, 100L, 25L, 75L), 500L, "Partial overlap interpolates");
+ assertEquals(StatsAggregations.overlapRows(100L, 50L, 50L, 0L, 100L), 100L, "Zero-length segment in range");
+ assertEquals(StatsAggregations.overlapRows(100L, 50L, 50L, 60L, 100L), 0L, "Zero-length segment out of range");
+ }
+
+ // ---------------------------------------------------------------------------
+ // Enumerating stored tables
+ // ---------------------------------------------------------------------------
+
+ /// getTables() feeds a destructive purge, so both stores must agree on what "holds statistics"
+ /// means -- including that a table whose segments are all gone is no longer reported.
+ @Test
+ public void testGetTablesReflectsStoredRows()
+ throws Exception {
+ assertTrue(_store.getTables().isEmpty(), "A fresh store holds nothing");
+
+ _store.upsertSegmentStats(TABLE_A, List.of(seg("s1", 1L, 100L, 1000L, 0L, 10L, false)));
+ _store.upsertSegmentStats(TABLE_B, List.of(seg("s2", 2L, 200L, 2000L, 0L, 10L, false)));
+ assertEquals(_store.getTables(), Set.of(TABLE_A, TABLE_B));
+
+ _store.removeSegments(TABLE_A, List.of("s1"));
+ assertEquals(_store.getTables(), Set.of(TABLE_B), "A table with no rows left is not reported");
+
+ _store.purgeTable(TABLE_B);
+ assertTrue(_store.getTables().isEmpty());
+ }
+
+ // ---------------------------------------------------------------------------
+ // Closed store
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testReadsAfterCloseAreRejected()
+ throws Exception {
+ _store.close();
+ assertThrows(StatsStoreException.class, () -> _store.getTableStats(TABLE_A));
+ assertThrows(StatsStoreException.class, () -> _store.getTables());
+ _store = null; // already closed; tearDown must not close it twice
+ }
+
+ // ---------------------------------------------------------------------------
+ // getSegmentCrcs
+ // ---------------------------------------------------------------------------
+
+ @Test
+ public void testGetSegmentCrcsEmptyWhenNoData()
+ throws Exception {
+ Map<String, Long> crcs = _store.getSegmentCrcs(TABLE_A);
+ assertNotNull(crcs);
+ assertTrue(crcs.isEmpty());
+ }
+
+ // ---------------------------------------------------------------------------
+ // Factory helpers
+ // ---------------------------------------------------------------------------
+
+ private static SegmentStatsRow seg(String name, long crc, long docs, long sizeBytes,
+ long startMs, long endMs, boolean consuming) {
+ return new SegmentStatsRow(name, crc, docs, sizeBytes, startMs, endMs, consuming);
+ }
+
+ /// Existing cases use numeric min/max, so they record LONG; the untyped case is exercised
+ /// explicitly by testUnknownValueTypeYieldsUntrustedBounds.
+ private static SegmentColumnStatsRow col(String segName, String colName, long ndv,
+ String minVal, String maxVal, boolean minTrusted, double avgBytes, double nullFrac) {
+ return new SegmentColumnStatsRow(segName, colName, ndv, minVal, maxVal, minTrusted, avgBytes,
+ nullFrac, ColumnValueType.LONG);
+ }
+}
diff --git a/pom.xml b/pom.xml
index 10eedc3643f..6b730e954a7 100644
--- a/pom.xml
+++ b/pom.xml
@@ -362,6 +362,7 @@
<equalsverifier.version>3.19.4</equalsverifier.version>
<testcontainers.version>2.0.5</testcontainers.version>
<h2.version>2.4.240</h2.version>
+ <sqlite-jdbc.version>3.46.1.3</sqlite-jdbc.version>
<jnr-posix.version>3.2.2</jnr-posix.version>
<scalatest.version>3.2.20</scalatest.version>
<assertj.version>3.27.7</assertj.version>
@@ -2084,6 +2085,13 @@
<scope>test</scope>
</dependency>
+ <!-- SQLite JDBC driver (broker-local stats store) -->
+ <dependency>
+ <groupId>org.xerial</groupId>
+ <artifactId>sqlite-jdbc</artifactId>
+ <version>${sqlite-jdbc.version}</version>
+ </dependency>
+
<!-- assertj bom -->
<dependency>
<groupId>org.assertj</groupId>
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org