feat(cubejs): enforce partition pruning in queryRewrite for cubes over partitioned event tables

services/cubejs/index.js

@@ -86,7 +86,7 @@ const options = {
   scheduledRefreshContexts,
   externalDbType: "cubestore",
   externalDriverFactory,
-  cacheAndQueueDriver: "cubestore",
+  cacheAndQueueDriver: process.env.CUBEJS_CACHE_AND_QUEUE_DRIVER || "cubestore",
   logger: logging,
 
   // sql server

services/cubejs/src/__tests__/partitionPruning.test.js

@@ -0,0 +1,117 @@
+import assert from "node:assert/strict";
+import { describe, it } from "node:test";
+
+import { applyPartitionPruning } from "../utils/queryRewrite.js";
+
+const policy = { semantic_events: { dimension: "timestamp", marginDays: 31 } };
+const map = new Map([
+  [
+    "simple_stays",
+    {
+      sourceTable: "semantic_events",
+      dimensions: new Set(["timestamp", "started", "poi"]),
+      partitionDimension: null,
+    },
+  ],
+  [
+    "bookings",
+    { sourceTable: "bookings", dimensions: new Set(["created"]), partitionDimension: null },
+  ],
+  [
+    "no_ts",
+    { sourceTable: "semantic_events", dimensions: new Set(["started"]), partitionDimension: null },
+  ],
+]);
+
+describe("partition pruning rewrite", () => {
+  it("adds a widened timestamp window from an explicit timeDimensions range", () => {
+    const q = applyPartitionPruning(
+      {
+        measures: ["simple_stays.poi_stay_count"],
+        dimensions: ["simple_stays.poi"],
+        timeDimensions: [
+          { dimension: "simple_stays.started", dateRange: ["2026-06-01", "2026-06-07"] },
+        ],
+      },
+      map,
+      policy,
+    );
+    assert.deepEqual(q.filters, [
+      {
+        member: "simple_stays.timestamp",
+        operator: "inDateRange",
+        values: ["2026-05-01", "2026-07-08"],
+      },
+    ]);
+  });
+
+  it("derives the window from inDateRange filters too, taking the union", () => {
+    const q = applyPartitionPruning(
+      {
+        measures: ["simple_stays.count"],
+        filters: [
+          { member: "simple_stays.started", operator: "inDateRange", values: ["2026-01-10", "2026-01-20"] },
+          { member: "simple_stays.started", operator: "inDateRange", values: ["2026-03-01", "2026-03-05"] },
+        ],
+      },
+      map,
+      policy,
+    );
+    const added = q.filters.at(-1);
+    assert.equal(added.member, "simple_stays.timestamp");
+    assert.deepEqual(added.values, ["2025-12-10", "2026-04-05"]);
+  });
+
+  it("leaves queries alone when the partition column is already constrained", () => {
+    const q = {
+      measures: ["simple_stays.count"],
+      timeDimensions: [
+        { dimension: "simple_stays.timestamp", dateRange: ["2026-06-01", "2026-06-07"] },
+      ],
+    };
+    applyPartitionPruning(q, map, policy);
+    assert.equal(q.filters, undefined);
+  });
+
+  it("skips relative ranges, ungoverned tables, and cubes without the dimension", () => {
+    for (const query of [
+      {
+        measures: ["simple_stays.count"],
+        timeDimensions: [{ dimension: "simple_stays.started", dateRange: "last 7 days" }],
+      },
+      {
+        measures: ["bookings.count"],
+        timeDimensions: [{ dimension: "bookings.created", dateRange: ["2026-06-01", "2026-06-07"] }],
+      },
+      {
+        measures: ["no_ts.count"],
+        timeDimensions: [{ dimension: "no_ts.started", dateRange: ["2026-06-01", "2026-06-07"] }],
+      },
+    ]) {
+      applyPartitionPruning(query, map, policy);
+      assert.equal(query.filters, undefined);
+    }
+  });
+
+  it("honours meta.partition_dimension per cube", () => {
+    const custom = new Map([
+      [
+        "events_v2",
+        {
+          sourceTable: "semantic_events",
+          dimensions: new Set(["event_time", "occurred"]),
+          partitionDimension: "event_time",
+        },
+      ],
+    ]);
+    const q = applyPartitionPruning(
+      {
+        measures: ["events_v2.count"],
+        timeDimensions: [{ dimension: "events_v2.occurred", dateRange: ["2026-02-01", "2026-02-02"] }],
+      },
+      custom,
+      policy,
+    );
+    assert.equal(q.filters[0].member, "events_v2.event_time");
+  });
+});

services/cubejs/src/utils/queryRewrite.js

@@ -11,6 +11,133 @@ const getColumnsArray = (cube) => [
   ...(cube?.segments || []),
 ];
 
+// --- Partition pruning policy (owner rule, 2026-09-03) ---
+// A cube over a partitioned event table must never reach the warehouse
+// without a predicate on the partition time column. `cst.semantic_events`
+// is partitioned by (partition, toStartOfMonth(timestamp)); a row-type cube
+// whose time dimension is a JSON payload path (e.g. `properties.started_at`)
+// makes ClickHouse read the JSON column for every row of every month —
+// measured 10 GiB / 18.6M rows for a one-week POI query, 13 MiB / 119 ms
+// once the timestamp window is present. The rewrite derives that window from
+// the query's explicit date ranges, widened by `marginDays` on both sides so
+// events recorded later than the payload time (a stay ends days after it
+// starts) are never excluded. Policy per source table; override with
+// CUBEJS_PARTITION_PRUNING (JSON: { table: { dimension, marginDays } }), and
+// per cube with `meta.partition_dimension`.
+const DEFAULT_PARTITION_PRUNING = {
+  semantic_events: { dimension: "timestamp", marginDays: 31 },
+};
+
+function loadPartitionPruningPolicy() {
+  const raw = process.env.CUBEJS_PARTITION_PRUNING;
+  if (!raw) return DEFAULT_PARTITION_PRUNING;
+  try {
+    const parsed = JSON.parse(raw);
+    return parsed && typeof parsed === "object"
+      ? { ...DEFAULT_PARTITION_PRUNING, ...parsed }
+      : DEFAULT_PARTITION_PRUNING;
+  } catch (err) {
+    console.error(
+      "[queryRewrite] CUBEJS_PARTITION_PRUNING is not valid JSON:",
+      err.message,
+    );
+    return DEFAULT_PARTITION_PRUNING;
+  }
+}
+
+const PARTITION_PRUNING = loadPartitionPruningPolicy();
+
+const DAY_MS = 24 * 60 * 60 * 1000;
+
+function parseExplicitDate(value) {
+  if (typeof value !== "string") return null;
+  // Cube accepts "YYYY-MM-DD" and ISO datetimes; relative expressions
+  // ("last 7 days") are left to Cube and yield no pruning window.
+  if (!/^\d{4}-\d{2}-\d{2}/.test(value)) return null;
+  const ms = Date.parse(value.length === 10 ? `${value}T00:00:00Z` : value);
+  return Number.isFinite(ms) ? ms : null;
+}
+
+function formatDay(ms) {
+  return new Date(ms).toISOString().slice(0, 10);
+}
+
+/**
+ * Explicit [from, to] windows the query states for members of `cubeName`,
+ * from timeDimensions.dateRange and inDateRange filters.
+ */
+function explicitDateWindows(query, cubeName) {
+  const windows = [];
+  const consider = (member, range) => {
+    if (typeof member !== "string" || !member.startsWith(`${cubeName}.`)) {
+      return;
+    }
+    const pair = Array.isArray(range) ? range : [range, range];
+    if (pair.length !== 2) return;
+    const from = parseExplicitDate(pair[0]);
+    const to = parseExplicitDate(pair[1]);
+    if (from === null || to === null) return;
+    windows.push([from, to]);
+  };
+  for (const td of query.timeDimensions || []) {
+    if (td?.dateRange) consider(td.dimension, td.dateRange);
+  }
+  for (const f of query.filters || []) {
+    if (f?.operator === "inDateRange") consider(f.member, f.values);
+  }
+  return windows;
+}
+
+function mentionsMember(query, member) {
+  return (
+    (query.timeDimensions || []).some((td) => td?.dimension === member) ||
+    (query.filters || []).some((f) => f?.member === member) ||
+    (query.dimensions || []).includes(member)
+  );
+}
+
+/**
+ * Add the partition-column window for every query cube whose source table
+ * is governed by the pruning policy. Mutates and returns `query`.
+ * Exported for tests.
+ */
+export function applyPartitionPruning(
+  query,
+  cubeToTable,
+  policy = PARTITION_PRUNING,
+) {
+  if (!query || !cubeToTable || cubeToTable.size === 0) return query;
+  if (process.env.CUBEJS_PARTITION_PRUNING_DEBUG === "1") {
+    console.log("[queryRewrite] pruning map", {
+      cubes: [...cubeToTable.entries()].map(([k, v]) => `${k}→${v.sourceTable}`),
+      query: extractCubeNames(query),
+    });
+  }
+  for (const cubeName of extractCubeNames(query)) {
+    const info = cubeToTable.get(cubeName);
+    if (!info?.sourceTable) continue;
+    const rule = policy[info.sourceTable];
+    if (!rule) continue;
+    const dimension = info.partitionDimension || rule.dimension;
+    if (!dimension || !info.dimensions?.has(dimension)) continue;
+    const member = `${cubeName}.${dimension}`;
+    // Already constrained on the partition column — nothing to add.
+    if (mentionsMember(query, member)) continue;
+    const windows = explicitDateWindows(query, cubeName);
+    if (windows.length === 0) continue;
+    const margin = (Number(rule.marginDays) || 0) * DAY_MS;
+    const from = Math.min(...windows.map((w) => w[0])) - margin;
+    const to = Math.max(...windows.map((w) => w[1])) + margin;
+    if (!query.filters) query.filters = [];
+    query.filters.push({
+      member,
+      operator: "inDateRange",
+      values: [formatDay(from), formatDay(to)],
+    });
+  }
+  return query;
+}
+
 // --- Rule cache with 60-second TTL ---
 let rulesCache = null;
 let rulesCacheTime = 0;
@@ -136,7 +263,11 @@ async function buildCubeToTableMap(schemaVersion, fileIds) {
         const dims = new Set(
           (cube.dimensions || []).map((d) => d.name).filter(Boolean),
         );
-        mapping.set(cube.name, { sourceTable, dimensions: dims });
+        mapping.set(cube.name, {
+          sourceTable,
+          dimensions: dims,
+          partitionDimension: cube.meta?.partition_dimension || null,
+        });
       }
     }
   } catch (err) {
@@ -214,16 +345,17 @@ const queryRewrite = async (query, { securityContext }) => {
     }
   }
 
+  // --- Step 0b: Partition pruning (applies to ALL roles, before any rule) ---
+  const { schemaVersion, files } = userScope.dataSource;
+  const cubeToTable = await buildCubeToTableMap(schemaVersion, files);
+  applyPartitionPruning(query, cubeToTable);
+
   // --- Step 1: Rule-based row filtering (applies to ALL roles) ---
   const rules = await loadRules();
 
   if (rules.length > 0) {
     const queryCubeNames = extractCubeNames(query);
 
-    // Build cube → source table mapping from active schemas
-    const { schemaVersion, files } = userScope.dataSource;
-    const cubeToTable = await buildCubeToTableMap(schemaVersion, files);
-
     // Index rules by table name for fast lookup
     const rulesByTable = new Map();
     for (const rule of rules) {