feat(adapter): route gauge grouped reads to AMT rollups

lohanidamodar · lohanidamodar · commit 720af37dd0f2 · 2026-06-18T03:41:04.000Z
diff --git a/src/Usage/Adapter/ClickHouse.php b/src/Usage/Adapter/ClickHouse.php
@@ -1896,7 +1896,7 @@ public function find(array $queries = [], ?string $type = null): array
         if ($type !== null) {
             if ($this->useDailyRollups && $type === Usage::TYPE_EVENT) {
                 $plan = $this->extractRoutingPlan($queries);
-                $route = $this->selectAggregateSource($plan);
+                $route = $this->selectAggregateSource($plan, Usage::TYPE_EVENT);
                 $this->recordRoute('find', $plan, $route);
 
                 if (str_starts_with($route, 'daily_by_')) {
@@ -1918,6 +1918,32 @@ public function find(array $queries = [], ?string $type = null): array
                     }
                 }
             }
+
+            if ($this->useDailyRollups && $type === Usage::TYPE_GAUGE) {
+                $plan = $this->extractRoutingPlan($queries);
+                $route = $this->selectAggregateSource($plan, Usage::TYPE_GAUGE);
+                $this->recordRoute('find', $plan, $route);
+
+                if (str_starts_with($route, 'gauges_daily_by_')) {
+                    $name = substr($route, strlen('gauges_daily_'));
+                    $rollup = $this->getGaugeRollupByName($name);
+                    if ($rollup !== null) {
+                        $result = $this->findFromGaugeDailyByDim($rollup, $queries);
+                        $this->maybeDualRead($queries, $type, $route, $plan, $result);
+                        return $result;
+                    }
+                }
+                if (str_starts_with($route, 'gauges_hybrid_by_')) {
+                    $name = substr($route, strlen('gauges_hybrid_'));
+                    $rollup = $this->getGaugeRollupByName($name);
+                    if ($rollup !== null) {
+                        $result = $this->findHybridFromGaugeDailyByDim($rollup, $queries);
+                        $this->maybeDualRead($queries, $type, $route, $plan, $result);
+                        return $result;
+                    }
+                }
+            }
+
             return $this->findFromTable($queries, $type);
         }
 
@@ -2434,22 +2460,25 @@ private function extractRoutingPlan(array $queries): array
      * Pure routing decision: which physical table satisfies this read.
      *
      * Returns one of:
-     *   - 'raw'    — scan the events table.
-     *   - 'daily'  — read the existing daily MV (no grouping, only
+     *   - 'raw'    — scan the raw events / gauges table.
+     *   - 'daily'  — read the existing events daily MV (no grouping, only
      *                daily-MV-compatible filters, window in closed days).
-     *   - 'hybrid' — closed days from daily MV, today's partial from raw.
+     *   - 'hybrid' — closed days from events daily MV, today's partial from raw.
      *   - 'daily_by_path' / 'daily_by_country' / 'daily_by_service' /
-     *     'daily_by_method_status' — read a per-dim rollup when the
-     *     requested dimensions[] and filter columns match one of the
-     *     four MVs.
-     *   - 'hybrid_by_<dim>' — multi-dim MV closed days + raw today's partial.
+     *     'daily_by_method_status' — read an events per-dim rollup.
+     *   - 'hybrid_by_<dim>' — events multi-dim MV closed days + raw today's
+     *     partial.
+     *   - 'gauges_daily_by_service' / 'gauges_daily_by_resource' — read a
+     *     gauges per-dim AMT rollup.
+     *   - 'gauges_hybrid_by_<dim>' — gauges multi-dim AMT MV closed days +
+     *     raw today's partial.
      *
      * Any caller-requested sub-day interval forces 'raw' because the
      * rollups are day-grained.
      *
      * @param array{metric: ?string, start: ?string, end: ?string, filterColumns: array<int, string>, dimensions: array<int, string>, interval: ?string} $plan
      */
-    private function selectAggregateSource(array $plan): string
+    private function selectAggregateSource(array $plan, string $type = Usage::TYPE_EVENT): string
     {
         if ($plan['interval'] !== null) {
             return 'raw';
@@ -2467,6 +2496,27 @@ private function selectAggregateSource(array $plan): string
         }
         $straddlesToday = $endDt >= $boundaryDt;
 
+        if ($type === Usage::TYPE_GAUGE) {
+            if (empty($plan['dimensions'])) {
+                return 'raw';
+            }
+
+            $dimMatch = $this->matchGaugeDimRollup($plan['dimensions']);
+            if ($dimMatch === null) {
+                return 'raw';
+            }
+
+            $allowedFilters = array_merge(['metric', 'time', 'tenant'], $dimMatch['dims']);
+            foreach ($plan['filterColumns'] as $column) {
+                if (!in_array($column, $allowedFilters, true)) {
+                    return 'raw';
+                }
+            }
+
+            $base = 'gauges_daily_' . $dimMatch['name'];
+            return $straddlesToday ? 'gauges_hybrid_' . $dimMatch['name'] : $base;
+        }
+
         if (empty($plan['dimensions'])) {
             foreach ($plan['filterColumns'] as $column) {
                 if (!in_array($column, self::DAILY_COLUMNS, true) && $column !== 'tenant') {
@@ -2493,8 +2543,8 @@ private function selectAggregateSource(array $plan): string
     }
 
     /**
-     * Find the per-dim rollup whose dim set matches the request exactly
-     * (order-insensitive). Returns null when no MV covers this shape.
+     * Find the events per-dim rollup whose dim set matches the request
+     * exactly (order-insensitive). Returns null when no MV covers this shape.
      *
      * @param array<int, string> $dimensions
      * @return array{name: string, dims: array<int, string>}|null
@@ -2514,6 +2564,41 @@ private function matchDimRollup(array $dimensions): ?array
         return null;
     }
 
+    /**
+     * Find the gauge per-dim AMT rollup whose dim set matches the request
+     * exactly (order-insensitive). Returns null when no MV covers this shape.
+     *
+     * @param array<int, string> $dimensions
+     * @return array{name: string, dims: array<int, string>}|null
+     */
+    private function matchGaugeDimRollup(array $dimensions): ?array
+    {
+        $needle = $dimensions;
+        sort($needle);
+
+        foreach (self::GAUGE_DIM_ROLLUPS as $rollup) {
+            $candidate = $rollup['dims'];
+            sort($candidate);
+            if ($candidate === $needle) {
+                return $rollup;
+            }
+        }
+        return null;
+    }
+
+    /**
+     * @return array{name: string, dims: array<int, string>}|null
+     */
+    private function getGaugeRollupByName(string $name): ?array
+    {
+        foreach (self::GAUGE_DIM_ROLLUPS as $rollup) {
+            if ($rollup['name'] === $name) {
+                return $rollup;
+            }
+        }
+        return null;
+    }
+
     private function stringifyTime(mixed $value): ?string
     {
         if ($value instanceof \DateTime) {
@@ -2674,6 +2759,116 @@ private function findHybridFromDailyByDim(array $rollup, array $queries): array
         return $this->parseAggregatedResults($result, Usage::TYPE_EVENT);
     }
 
+    /**
+     * Day-grained aggregated read from a per-dim gauge AMT rollup table.
+     * Same return shape as findFromDailyByDim but emits argMaxMerge(value)
+     * instead of sum(value) — gauges store argMaxState partials.
+     *
+     * @param array{name: string, dims: array<int, string>} $rollup
+     * @param array<Query> $queries
+     * @return array<Metric>
+     */
+    private function findFromGaugeDailyByDim(array $rollup, array $queries): array
+    {
+        $tableName = $this->getGaugeDimRollupTableName($rollup['name']);
+        $fromTable = $this->buildTableReference($tableName);
+
+        $parsed = $this->parseQueries($queries, Usage::TYPE_GAUGE);
+        $whereData = $this->buildWhereClause($parsed['filters'], $parsed['params']);
+
+        $dimSelect = '';
+        $dimGroup = '';
+        if (!empty($rollup['dims'])) {
+            $escapedDims = array_map(fn (string $d): string => $this->escapeIdentifier($d), $rollup['dims']);
+            $dimSelect = ', ' . implode(', ', $escapedDims);
+            $dimGroup = ', ' . implode(', ', $escapedDims);
+        }
+
+        $orderClause = ' ORDER BY value DESC';
+        if (!empty($parsed['orderBy'])) {
+            $orderClause = ' ORDER BY ' . implode(', ', $parsed['orderBy']);
+        }
+        $limitClause = isset($parsed['limit']) ? ' LIMIT {limit:UInt64}' : '';
+        $offsetClause = isset($parsed['offset']) ? ' OFFSET {offset:UInt64}' : '';
+
+        $sql = "
+            SELECT metric, argMaxMerge(value) AS value{$dimSelect}
+            FROM {$fromTable}{$whereData['clause']}
+            GROUP BY metric{$dimGroup}{$orderClause}{$limitClause}{$offsetClause}
+            FORMAT JSON
+        ";
+
+        $result = $this->query($sql, $whereData['params']);
+        return $this->parseAggregatedResults($result, Usage::TYPE_GAUGE);
+    }
+
+    /**
+     * Hybrid gauge AMT read: closed-day rollup partials (argMaxState) UNION
+     * raw gauges (argMaxState(value, time) over today's partial), with an
+     * outer argMaxMerge that combines both sides. Same return shape as
+     * findHybridFromDailyByDim but argMax-based.
+     *
+     * @param array{name: string, dims: array<int, string>} $rollup
+     * @param array<Query> $queries
+     * @return array<Metric>
+     */
+    private function findHybridFromGaugeDailyByDim(array $rollup, array $queries): array
+    {
+        $startOfToday = (new \DateTime('today', new \DateTimeZone('UTC')))->format('Y-m-d H:i:s.v');
+
+        $dailyName = $this->getGaugeDimRollupTableName($rollup['name']);
+        $dailyTable = $this->buildTableReference($dailyName);
+        $gaugesTable = $this->buildTableReference($this->getGaugesTableName());
+
+        $parsed = $this->parseQueries($queries, Usage::TYPE_GAUGE);
+        $rawFilters = $parsed['filters'];
+        $params = $parsed['params'];
+
+        $dailyFilters = $rawFilters;
+
+        $params['hybrid_boundary'] = $startOfToday;
+        $dailyFilters[] = '`time` < {hybrid_boundary:DateTime64(3)}';
+        $rawFilters[] = '`time` >= {hybrid_boundary:DateTime64(3)}';
+
+        $dailyWhere = $this->buildWhereClause($dailyFilters, $params);
+        $rawWhere = $this->buildWhereClause($rawFilters, $dailyWhere['params']);
+
+        $dimSelect = '';
+        $dimGroup = '';
+        $dimCols = '';
+        if (!empty($rollup['dims'])) {
+            $escapedDims = array_map(fn (string $d): string => $this->escapeIdentifier($d), $rollup['dims']);
+            $dimList = implode(', ', $escapedDims);
+            $dimSelect = ', ' . $dimList;
+            $dimGroup = ', ' . $dimList;
+            $dimCols = ', ' . $dimList;
+        }
+
+        $orderClause = ' ORDER BY value DESC';
+        if (!empty($parsed['orderBy'])) {
+            $orderClause = ' ORDER BY ' . implode(', ', $parsed['orderBy']);
+        }
+        $limitClause = isset($parsed['limit']) ? ' LIMIT {limit:UInt64}' : '';
+        $offsetClause = isset($parsed['offset']) ? ' OFFSET {offset:UInt64}' : '';
+
+        $sql = "
+            SELECT metric, argMaxMerge(value) AS value{$dimSelect}
+            FROM (
+                SELECT metric{$dimCols}, value
+                FROM {$dailyTable}{$dailyWhere['clause']}
+                UNION ALL
+                SELECT metric{$dimCols}, argMaxState(value, time) AS value
+                FROM {$gaugesTable}{$rawWhere['clause']}
+                GROUP BY metric{$dimGroup}
+            )
+            GROUP BY metric{$dimGroup}{$orderClause}{$limitClause}{$offsetClause}
+            FORMAT JSON
+        ";
+
+        $result = $this->query($sql, $rawWhere['params']);
+        return $this->parseAggregatedResults($result, Usage::TYPE_GAUGE);
+    }
+
     /**
      * Dual-read sampler: with probability `$dualReadSampleRate`, re-run
      * the same logical query against raw events and log a warning if the
diff --git a/tests/Benchmark/BenchmarkBase.php b/tests/Benchmark/BenchmarkBase.php
@@ -130,10 +130,16 @@ protected function seedGaugeRows(int $rows, string $metric = 'storage'): void
 
         $tableRef = "`{$database}`.`{$gaugesTable}`";
 
-        $sql = "INSERT INTO {$tableRef} (id, metric, value, time, tenant) "
+        $services = ['storage', 'databases', 'functions', 'sites'];
+        $resources = ['file', 'database', 'function', 'site'];
+        $svcExpr = "['" . implode("','", $services) . "'][1 + (number % " . count($services) . ")]";
+        $resExpr = "['" . implode("','", $resources) . "'][1 + (number % " . count($resources) . ")]";
+
+        $sql = "INSERT INTO {$tableRef} (id, metric, value, time, tenant, service, resource) "
             . "SELECT lower(hex(randomString(16))), '" . addslashes($metric) . "', "
             . "number AS value, now() - toIntervalSecond(number % 86400) AS time, "
-            . "'" . addslashes($this->tenant) . "' AS tenant "
+            . "'" . addslashes($this->tenant) . "' AS tenant, "
+            . "{$svcExpr} AS service, {$resExpr} AS resource "
             . "FROM numbers({$rows})";
 
         $this->runRawSql($sql);
diff --git a/tests/Benchmark/GaugesBench.php b/tests/Benchmark/GaugesBench.php
@@ -4,6 +4,7 @@
 
 use Utopia\Query\Query;
 use Utopia\Usage\Usage;
+use Utopia\Usage\UsageQuery;
 
 class GaugesBench extends BenchmarkBase
 {
@@ -20,21 +21,59 @@ protected function setUp(): void
 
     public function testBenchmarks(): void
     {
-        $start = (new \DateTime('-30 days'))->format('Y-m-d H:i:s');
-        $end = (new \DateTime('+1 day'))->format('Y-m-d H:i:s');
+        $start30d = (new \DateTime('-30 days'))->format('Y-m-d H:i:s');
+        $endClosed = (new \DateTime('-2 days'))->format('Y-m-d H:i:s');
+        $endPartial = (new \DateTime('+1 day'))->format('Y-m-d H:i:s');
 
-        $this->runBench('bench_gauges_latest_in_window', function (string $queryId) use ($start, $end): void {
+        $this->runBench('bench_gauges_latest_in_window', function (string $queryId) use ($start30d, $endPartial): void {
             $this->adapter->setNextQueryId($queryId);
             $this->usage->getTotal(
                 $this->metric,
                 [
-                    Query::greaterThanEqual('time', $start),
-                    Query::lessThanEqual('time', $end),
+                    Query::greaterThanEqual('time', $start30d),
+                    Query::lessThanEqual('time', $endPartial),
                 ],
                 Usage::TYPE_GAUGE
             );
         });
 
+        $this->adapter->setUseDailyRollups(true);
+
+        $this->runBench('bench_gauges_topN_service_30d', function (string $queryId) use ($start30d, $endClosed): void {
+            $this->adapter->setNextQueryId($queryId);
+            $this->usage->find([
+                UsageQuery::groupBy('service'),
+                Query::equal('metric', [$this->metric]),
+                Query::greaterThanEqual('time', $start30d),
+                Query::lessThanEqual('time', $endClosed),
+                Query::limit(10),
+            ], Usage::TYPE_GAUGE);
+        });
+
+        $this->runBench('bench_gauges_topN_resource_30d', function (string $queryId) use ($start30d, $endClosed): void {
+            $this->adapter->setNextQueryId($queryId);
+            $this->usage->find([
+                UsageQuery::groupBy('resource'),
+                Query::equal('metric', [$this->metric]),
+                Query::greaterThanEqual('time', $start30d),
+                Query::lessThanEqual('time', $endClosed),
+                Query::limit(10),
+            ], Usage::TYPE_GAUGE);
+        });
+
+        $this->runBench('bench_gauges_topN_service_today_partial', function (string $queryId) use ($start30d, $endPartial): void {
+            $this->adapter->setNextQueryId($queryId);
+            $this->usage->find([
+                UsageQuery::groupBy('service'),
+                Query::equal('metric', [$this->metric]),
+                Query::greaterThanEqual('time', $start30d),
+                Query::lessThanEqual('time', $endPartial),
+                Query::limit(10),
+            ], Usage::TYPE_GAUGE);
+        });
+
+        $this->adapter->setUseDailyRollups(false);
+
         $this->assertNotEmpty($this->results, 'Benchmark scenarios must record results');
     }
 }
diff --git a/tests/Usage/Adapter/ClickHouseGaugeDimRoutingTest.php b/tests/Usage/Adapter/ClickHouseGaugeDimRoutingTest.php
