From 55a727cb1e88d03e84d56b1ccb037f39efd5a120 Mon Sep 17 00:00:00 2001 From: stefanbaxter Date: Thu, 3 Sep 2026 05:39:36 +0000 Subject: [PATCH 1/2] feat(cubejs): enforce partition pruning in queryRewrite for cubes over partitioned event tables Owner rule (cxs2 spec 107 follow-up, 2026-09-03): a cube over cst.semantic_events must never reach ClickHouse without a predicate on the partition time column. A row-type cube whose time dimension is a payload path (simple_stays.started = properties.started_at) made every dashboard query read the JSON column across all monthly partitions: 10 GiB / 18.6M rows for a one-week POI query, killed by the server memory guard whenever two ran together. queryRewrite now derives a window from the query's explicit date ranges (timeDimensions.dateRange, inDateRange filters) and adds an inDateRange filter on the cube's partition dimension, widened by a margin on both sides. Policy per source table (default: semantic_events -> timestamp, 31 days), overridable with CUBEJS_PARTITION_PRUNING and per cube with meta.partition_dimension. It reuses the existing cube-to-table map, runs for every role, and stays out of the way when the query already constrains the partition column or states no explicit range. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Nv1ngm1j8mV8rWsumTbpx2 --- .../src/__tests__/partitionPruning.test.js | 117 +++++++++++++++ services/cubejs/src/utils/queryRewrite.js | 142 +++++++++++++++++- 2 files changed, 254 insertions(+), 5 deletions(-) create mode 100644 services/cubejs/src/__tests__/partitionPruning.test.js diff --git a/services/cubejs/src/__tests__/partitionPruning.test.js b/services/cubejs/src/__tests__/partitionPruning.test.js new file mode 100644 index 00000000..d91991dd --- /dev/null +++ b/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"); + }); +}); diff --git a/services/cubejs/src/utils/queryRewrite.js b/services/cubejs/src/utils/queryRewrite.js index 34c0beb1..84e42e4a 100644 --- a/services/cubejs/src/utils/queryRewrite.js +++ b/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) { From 0d06cf4c491dcaa9d6c91b955c96855e1529e6e8 Mon Sep 17 00:00:00 2001 From: stefanbaxter Date: Thu, 3 Sep 2026 05:51:30 +0000 Subject: [PATCH 2/2] chore(cubejs): let CUBEJS_CACHE_AND_QUEUE_DRIVER override the cubestore default Local dev on Apple silicon runs cubestore under amd64 emulation, where the cache driver's WebSocket dies at startup and every query fails with 'Cube Store connection is closed'. Honouring the standard Cube env var (default unchanged: cubestore) lets the dev stack run with the in-memory driver; production configs set nothing and keep cubestore. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01Nv1ngm1j8mV8rWsumTbpx2 --- services/cubejs/index.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/services/cubejs/index.js b/services/cubejs/index.js index 5a09747d..79524aa0 100644 --- a/services/cubejs/index.js +++ b/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