[OPIK-6901] [BE] perf: prune trace deletes to the batch's own partitions

apps/opik-backend/config.yml

@@ -155,6 +155,20 @@ databaseAnalyticsDataModel:
   #   MODIFY TTL target `traces_local` only (the Distributed `traces` rejects them); ADD/DROP/MODIFY COLUMN must target
   #   both `traces_local` and `traces`, else reads can't see the column (code 47).
   tracesDistributedWrapEnabled: ${ANALYTICS_DB_DATA_MODEL_TRACES_DISTRIBUTED_WRAP_ENABLED:-false}
+  #  Default: false
+  #  Description: Enables partition-aware PRUNING of trace deletes - it does NOT create or activate any partitioning.
+  #   With it on, a trace DELETE bounds itself to the weekly partitions its own ids resolve to instead of being planned
+  #   against every part of the table. Turning it on therefore ASSERTS a schema fact rather than causing one: that the
+  #   live mutation target already IS the weekly partitioned successor, id_at as DateTime64(0,'UTC') under
+  #   PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))). Installing that schema is the
+  #   EXCHANGE step of the cutover, never this flag. Purely an optimisation: false keeps the unbounded mutation, which
+  #   is always correct and merely slower. A third flag on purpose - the partitioning appears at the EXCHANGE, and
+  #   neither sibling marks it:
+  #   traceColumnsNonNullable must be rolled out BEFORE the EXCHANGE, tracesDistributedWrapEnabled flips at the wrap,
+  #   which may be deferred long after it. Leave false at deploy time; set true once the EXCHANGE is confirmed, and back
+  #   to false BEFORE a rollback promotes the original `traces` (legacy `traces` has no PARTITION BY and a 32-bit
+  #   DateTime id_at that overflows past 2106, so the predicate would silently match zero rows for a far-future id).
+  tracesWeeklyPartitionPruningEnabled: ${ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED:-false}
 
 # Description: UUIDv7 ingestion validation. Rejects writes whose `id` embeds a timestamp outside the
 #   window, protecting data quality.

apps/opik-backend/data-migrations/traces-local-v2-cutover/README.md

@@ -322,6 +322,30 @@ the `traceColumnsNonNullable` flip").
 
 On rollback, after swapping the Nullable original back, revert the flag to `false` **and** run that repair.
 
+**The `tracesWeeklyPartitionPruningEnabled` flip (optional, and why it goes last).** `databaseAnalyticsDataModel.tracesWeeklyPartitionPruningEnabled`
+(env `ANALYTICS_DB_DATA_MODEL_TRACES_WEEKLY_PARTITION_PRUNING_ENABLED`, default `false`) lets a trace `DELETE` bind itself to
+the weekly partitions its own ids resolve to (OPIK-6901), instead of being planned against every part of the table — on
+prod-test, 12 ids rewrote 3,928 parts / 5.40 TiB without it. It asserts a **schema** fact: that `traces` (or
+`traces_local`) is the successor, with `id_at` as `DateTime64(0,'UTC')` under the weekly `PARTITION BY`.
+
+> **It enables the *pruning*, not the partitioning — the name is deliberate.** Setting it does not create, activate or
+> migrate anything; the partitioned schema arrives with the `EXCHANGE` above and nowhere else. So it is never a step that
+> *makes* the cutover progress, and setting it early does not bring the partitioning forward — it only starts emitting a
+> predicate against whatever table is live, which is the failure below.
+
+It is a third flag precisely because **neither of the two above marks the `EXCHANGE`**, which is when the partitioning
+appears: `traceColumnsNonNullable` must lead it (above), and `tracesDistributedWrapEnabled` flips at the wrap, which may
+be deferred long after it (`--skip-wrap` … `--wrap-only`). So gate on this one, not on either of those.
+
+Unlike its siblings it is **safe to lag and unsafe to lead**: `false` is the previous unbounded delete, always correct
+and merely slower, so turn it on at leisure **after** the `EXCHANGE` is confirmed. Turning it on early — while `traces` is
+still the original — is the failure mode worth avoiding: the original has **no `PARTITION BY` at all** (nothing to prune)
+and declares `id_at` as a 32-bit `DateTime` that overflows past 2106, so a far-future id (the litellm ~2201 rows) is
+stored under a wrapped recent timestamp that the derived partition cannot match, and the delete reports success having
+matched **zero rows**. For the same reason, a stage B/C **rollback must revert it to `false` — and roll-restart every
+instance — before promoting the original**, ahead of the swap rather than after it. Like its siblings it comes from a
+startup snapshot of `OpikConfiguration`, so the config change alone changes nothing until each instance restarts.
+
 ## Batching and throttling
 
 On a large production table a single week can be enormous, so the backfill does **not** run one INSERT per week. Two
@@ -832,6 +856,11 @@ statements, so a failure *between* them needs a restart path:
 again, so the flip has to be undone in two steps — `rollback.sh` prints both when the stage finishes. The rollback is not
 complete until they land.
 
+> **If `tracesWeeklyPartitionPruningEnabled` was turned on, revert it *before* the stage runs, not after.** It asserts the
+> live table is the partitioned successor, and the restored original is not one, so a stale `true` makes trace deletes
+> match zero rows while reporting success — see "The `tracesWeeklyPartitionPruningEnabled` flip". It is the one flag whose
+> revert has to lead the swap; the two steps below follow it.
+
 1. **Revert `traceColumnsNonNullable` to `false` AND roll-restart every backend instance.** The flag is read from a
    **startup snapshot** of `OpikConfiguration` (bound via `toInstance`), so a config change does **not** take effect until
    each instance restarts — exactly like the forward rollout before the EXCHANGE. Until the restart completes, the app

apps/opik-backend/src/main/java/com/comet/opik/domain/TraceDAO.java

@@ -32,6 +32,7 @@
 import com.comet.opik.utils.ErrorUtils;
 import com.comet.opik.utils.JsonUtils;
 import com.comet.opik.utils.TruncationUtils;
+import com.comet.opik.utils.WeeklyPartitions;
 import com.comet.opik.utils.template.TemplateUtils;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.JsonNodeFactory;
@@ -1921,8 +1922,21 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC
      * so a single statement can span several projects (e.g. a reused id resolved to all its owning projects, or a
      * cross-project batch) instead of one delete per project (OPIK-7483). Every deleted row carries its {@code
      * project_id}, so no delete is ever project-less - also required once {@code traces} is a Distributed table
-     * (OPIK-7455). No {@code id_at}/time predicate on purpose, so it still deletes rows whose {@code id_at} is
-     * untrustworthy (e.g. a wrapped timestamp); correctness here does not depend on {@code id_at}.
+     * (OPIK-7455).
+     * <p>
+     * {@code <if(partitions)>} adds the table's own weekly partition expression, bound as the exact set of partitions
+     * the batch's ids resolve to. Both conditions must hold for it to be emitted: the live table must be the weekly
+     * partitioned successor ({@link #tracesWeeklyPartitionPruningEnabled()}), and every id in the batch must be one whose
+     * partition can be derived exactly ({@link WeeklyPartitions#of}). Otherwise the predicate is omitted and the
+     * statement is byte-identical to the previous unbounded form. That is what preserves the original guarantee — a
+     * row whose {@code id_at} cannot be trusted is still deleted, because no id in such a batch is used to derive a
+     * partition.
+     * <p>
+     * Why it matters: a mutation selects parts at the <b>partition</b> stage, where the (workspace_id, project_id, id)
+     * predicate prunes nothing, so deleting a handful of rows rewrote every part of the table. Measured on prod-test
+     * (271.6 M rows, 3,928 parts): 12 ids rewrote <b>3,928 parts / 5.40 TiB</b>. With this predicate the same batch
+     * selects <b>5</b> parts. An {@code id_at} <em>range</em> is not a substitute: on a batch spanning 1996 and 2200 a
+     * range still selected 2,644 parts, where the exact set selected 4.
      * <p>
      * The pairs are bound (never inlined) as two positional string arrays and zipped back into {@code (project_id, id)}
      * tuples with {@code arrayZip}, so the query text is constant regardless of batch size and no value reaches the SQL
@@ -1934,6 +1948,7 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC
             DELETE FROM <if(distributed_wrap)>traces_local<else>traces<endif>
             WHERE workspace_id = :workspace_id
             AND (project_id, id) IN arrayZip(:project_ids, :trace_ids)
+            <if(partitions)>AND toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1))) IN :partitions<endif>
             SETTINGS log_comment = '<log_comment>'
             ;
             """;
@@ -3356,6 +3371,49 @@ private void selectTracesMutationTable(ST template) {
         }
     }
 
+    /**
+     * Whether a trace mutation may prune to the partitions its ids resolve to ({@link WeeklyPartitions}). The flag
+     * enables the <b>pruning</b>, never the partitioning: it asserts that the live mutation target already is the
+     * weekly partitioned successor — {@code id_at} as {@code DateTime64(0, 'UTC')} under
+     * {@code PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))} — which the cutover's
+     * EXCHANGE installs, not this flag.
+     * <p>
+     * This has to be its own flag because <b>neither existing schema flag marks the EXCHANGE, which is the moment the
+     * partitioning appears</b>, and both are wrong in a different direction:
+     * <ul>
+     *     <li>{@link #tracesDistributedWrapEnabled()} is too late. The wrap is a separate, deferrable step after the
+     *     EXCHANGE (it may be skipped entirely with {@code --skip-wrap} and applied weeks later with
+     *     {@code --wrap-only}), so between the two {@code traces} is already the partitioned successor while the flag is
+     *     still {@code false} — the state prod-test sat in. Gating on it would simply forgo the pruning there.</li>
+     *     <li>{@link #traceColumnsNonNullable()} is too early, which is the dangerous direction. It is a runtime
+     *     concern, not a schema one, and the runbook requires it rolled out to {@code true} on every instance
+     *     <b>before</b> the EXCHANGE (a rolling restart cannot be atomic with a metadata swap). Gating on it would emit
+     *     the predicate against the legacy {@code traces} for the whole rollout window.</li>
+     * </ul>
+     * Emitting it against the legacy table is not merely unhelpful, it is wrong: legacy {@code traces} has no
+     * {@code PARTITION BY} at all (one {@code all} partition, so nothing to prune) and declares {@code id_at} as a
+     * 32-bit {@code DateTime} that overflows past 2106, so a far-future id — the litellm ~2201 ids, real
+     * customer-facing rows — is stored under a wrapped recent timestamp that the derived partition cannot match. The
+     * delete would then match zero rows and report success.
+     * <p>
+     * Only ever {@code true} while {@code traces} really is that successor, so unlike its siblings it is safe to lag:
+     * {@code false} is the always-correct unbounded behaviour, and only {@code true} asserts something about the
+     * schema. Turn it on once the EXCHANGE is confirmed, and back off <b>before</b> a rollback stage B/C promotes the
+     * original.
+     */
+    private boolean tracesWeeklyPartitionPruningEnabled() {
+        return configuration.getDatabaseAnalyticsDataModel().tracesWeeklyPartitionPruningEnabled();
+    }
+
+    /**
+     * The partitions a delete batch may bound itself to, or empty to leave the mutation unbounded. Empty whenever the
+     * live table is not the partitioned successor, ahead of asking {@link WeeklyPartitions} at all — the derivation is
+     * only meaningful against a table that partitions on it.
+     */
+    private Optional<Set<Long>> weeklyPartitionsFor(Collection<UUID> ids) {
+        return tracesWeeklyPartitionPruningEnabled() ? WeeklyPartitions.of(ids) : Optional.empty();
+    }
+
     /**
      * Binds input, output, metadata, and their slim versions (input_slim, output_slim) to a statement.
      * Centralizes the JSON conversion and binding logic for consistency across single and batch inserts.
@@ -3529,11 +3587,22 @@ public Mono<Void> delete(Set<Pair<UUID, UUID>> projectIdTraceIdPairs, @NonNull C
                     var projectIds = batch.stream().map(pair -> pair.getLeft().toString()).toArray(String[]::new);
                     var traceIds = batch.stream().map(pair -> pair.getRight().toString()).toArray(String[]::new);
 
+                    // Prune to the batch's own partitions when the schema and every id in the batch allow it;
+                    // otherwise emit the unbounded form.
+                    var partitions = weeklyPartitionsFor(batch.stream().map(Pair::getRight).toList());
+                    // Flag only, exactly like distributed_wrap: the values reach ClickHouse via the bind below,
+                    // never through the template, so the rendered SQL is constant regardless of batch contents.
+                    partitions.ifPresent(_ -> template.add("partitions", true));
+
                     var statement = connection.createStatement(template.render())
                             .bind("workspace_id", workspaceId)
                             .bind("project_ids", projectIds)
                             .bind("trace_ids", traceIds);
 
+                    if (partitions.isPresent()) {
+                        statement = statement.bind("partitions", partitions.get().toArray(Long[]::new));
+                    }
+
                     var segment = startSegment("traces", "Clickhouse", "delete");
                     return Mono.from(statement.execute())
                             .doFinally(_ -> endSegment(segment))

apps/opik-backend/src/main/java/com/comet/opik/infrastructure/DatabaseAnalyticsDataModelConfig.java

@@ -50,6 +50,26 @@
  * {@code MODIFY COLUMN} must be applied to <b>both</b> {@code traces_local} and the {@code Distributed} {@code traces}
  * (the wrapper accepts them as metadata-only, and targeting only {@code traces_local} leaves the wrapper without the
  * column, so reads fail with code 47).</p>
+ *
+ * <p>{@code tracesWeeklyPartitionPruningEnabled}: enables partition-aware <b>pruning</b> of trace deletes — a trace
+ * {@code DELETE} bounds itself to the weekly partitions its own ids resolve to instead of being planned against every
+ * part of the table (OPIK-6901). It does <b>not</b> create or activate any partitioning; installing the partitioned
+ * schema is the EXCHANGE step of the cutover. Turning it on therefore <b>asserts</b> a schema fact rather than causing
+ * one: that the live mutation target already is the weekly partitioned successor — {@code id_at} as
+ * {@code DateTime64(0, 'UTC')} under
+ * {@code PARTITION BY toYYYYMMDD(toDate32(id_at) - toIntervalDay(toDayOfWeek(id_at, 1)))}. Purely an optimisation:
+ * {@code false} keeps the unbounded mutation, which is always correct and merely slower, and only {@code true} asserts
+ * anything about the schema.</p>
+ *
+ * <p>It is deliberately a third flag rather than a reuse of the two above, because the partitioning appears at the
+ * <b>EXCHANGE</b> and neither of them marks that moment. {@code traceColumnsNonNullable} must be rolled out
+ * <b>before</b> the EXCHANGE (a rolling restart cannot be atomic with a metadata swap), and
+ * {@code tracesDistributedWrapEnabled} flips at the wrap, a separate step that may be deferred long after it — so one
+ * flag would be true too early and the other true too late. Emitting the predicate too early is the harmful direction:
+ * legacy {@code traces} has no {@code PARTITION BY} at all and declares {@code id_at} as a 32-bit {@code DateTime} that
+ * overflows past 2106, so a far-future id is stored under a wrapped timestamp the derived partition cannot match and
+ * the delete would silently affect zero rows. Left {@code false} at deploy time; set {@code true} once the EXCHANGE is
+ * confirmed on the target, and back to {@code false} <b>before</b> a rollback promotes the original {@code traces}.</p>
  */
 @Builder(toBuilder = true)
 public record DatabaseAnalyticsDataModelConfig(
@@ -58,5 +78,6 @@ public record DatabaseAnalyticsDataModelConfig(
         boolean traceDeletionEventsCaptureEnabled,
         boolean spanDeletionEventsCaptureEnabled,
         @Min(1) @Max(2_000) int deletionEventsInsertBatchSize,
-        boolean tracesDistributedWrapEnabled) {
+        boolean tracesDistributedWrapEnabled,
+        boolean tracesWeeklyPartitionPruningEnabled) {
 }