[OPIK-8230] fix(traces): scope the delete mutation to the batch's own partitions (#8131)

A ClickHouse mutation is registered against every active part of the target
table before its WHERE clause is even considered - the existing partition
predicate (OPIK_6901) prunes which rows a matched part rewrites, not which
parts the mutation visits at all. Post-cutover traces_local has ~1,900
partitions / ~3,650 parts, so an ordinary delete matching a handful of rows
was registered against all of them, at ~19ms/part fixed overhead. That is
what pushed deletes past max_execution_time (Code 159 TIMEOUT_EXCEEDED),
where the statement errors and skips the TracesDeleted cascade and the
deletion-events bridge write even though the mutation completes and the
rows are in fact deleted.

Scope each statement to the partitions its own batch resolves to via
IN PARTITION, derived by WeeklyPartitions#groupByPartition (new; groups
ids by partition rather than naming their union - IN PARTITION accepts
exactly one partition per statement). One statement per partition,
falling back to the unbounded form when the batch cannot be derived
exactly, or when the target table is not partitioned.

That last condition has to be gated, not inferred from the table name.
Unlike the id_at predicate it supersedes - MATERIALIZED on both topologies,
so valid everywhere and merely unpruned on an unpartitioned table -
IN PARTITION against a table with no PARTITION BY is a hard Code 248
INVALID_PARTITION_VALUE. A name check is wrong in both directions:
post-EXCHANGE/pre-wrap `traces` is already the partitioned successor while
the wrap flag is still false, so the check resolves to "traces", takes the
fallback on every delete, and does nothing in exactly the state production
is in today; and wrap-on without the EXCHANGE is a real supported topology
(TracesDistributedWrapMutationTest) where traces_local is the legacy
unpartitioned table and IN PARTITION fails outright.

The gate reuses traceColumnsNonNullable rather than adding a flag. The
EXCHANGE that replaces the Nullable(...) columns is the same one that puts
the weekly-partitioned successor behind the name mutations target, so the
two facts have only ever flipped together, and a second flag would have to
be carried through the cutover runbook and tooling to track no independent
state. Verified against all four topologies: legacy pre-cutover and
wrap-without-EXCHANGE both read false and take the unbounded form;
production today and the pruning suite both read true and scope. The name
says only its first duty and is deliberately not renamed - the env var is
exposed - so both it and the config comment now document the second.

This leaves no deployment surface: no new config key, no Helm or
docker-compose change, nothing to add to the runbook. config.yml and
config-test.yml carry comment-only edits documenting the flag's second
duty where an operator will actually look for it.

Both forms share one SQL template with an <if(partition)> conditional, and
WeeklyPartitions#of (the flat-union predecessor) is removed - it had no
production callers once TraceDAO switched to groupByPartition. Its unique
boundary-case coverage moved to groupByPartition rather than being dropped.

A null project or trace id in the batch is now rejected at the DAO's entry
point rather than surfacing as an NPE while stringifying the binds. The NPE
is pre-existing - `pair.getRight().toString()` is unchanged by this diff -
but the partitioned path makes it worse: WeeklyPartitions reads a null as
"underivable", so the batch silently takes the unbounded fallback, and a
caller's bug reports as the slow delete this change exists to remove. The
check sits before the gate, so it holds whichever branch deleteBatch takes.

Verified locally against real ClickHouse via testcontainers, 35/35:
- WeeklyPartitionsTest 17, TracesPartitionPruningMutationTest 9,
  TracesDistributedWrapMutationTest 4, TracesLegacyTablePruningMutationTest 2,
  TraceMutationRoutingArchTest 3
- The pruning suite asserts on system.part_log MutatePart events rather than
  EXPLAIN, which reports read-planner pruning - a different layer from what a
  mutation is registered against, and the gap the previous coverage fell into.
- Gate verified by negative control: forced off, the emitted SQL carries no
  IN PARTITION and 5 of the 9 pruning tests fail, so the gate is not vacuous.
- Null precondition likewise: removed, the new test reproduces exactly the
  reported NPE rather than the IllegalArgumentException it asserts.
- Log reads poll a fast-flushing container config instead of SYSTEM FLUSH
  LOGS, which pushed only what had already buffered and so raced the rows it
  was meant to reveal.

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: thiagohora

apps/opik-backend/config.yml

@@ -122,6 +122,13 @@ databaseAnalyticsDataModel:
   #   non-nullable columns so writes bind the sentinels (end_time->epoch, ttft->NaN) instead. Gates reads too, so
   #   while false a legitimate epoch end time round-trips unchanged instead of reading as null.
   #   Flip in lockstep with the cutover EXCHANGE.
+  #   Also gates partition-scoped trace deletes: that same EXCHANGE puts the weekly-partitioned successor behind the
+  #   name mutations target, so true additionally means a delete may scope itself with IN PARTITION instead of being
+  #   registered against every part of the table. The two facts have only ever flipped together, which is why this
+  #   gates both rather than there being a second flag. IN PARTITION against an unpartitioned table is a hard error
+  #   (code 248), not a lost optimisation, so setting this true ahead of the EXCHANGE breaks deletes - as it already
+  #   breaks writes, which bind sentinels a Nullable column has no use for. Off falls back to the unbounded delete:
+  #   correct, merely slower.
   traceColumnsNonNullable: ${ANALYTICS_DB_DATA_MODEL_TRACE_COLUMNS_NON_NULLABLE:-false}
   #  Default: false
   #  Description: Spans sibling of traceColumnsNonNullable. Leave false while the spans table still has Nullable

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

@@ -1967,40 +1967,22 @@ ORDER BY (workspace_id, project_id, trace_id, id) DESC, last_updated_at DESC
             """;
 
     /**
-     * Deletes by the full {@code (workspace_id, project_id, id)} sort key, matching on {@code (project_id, id)} tuples
-     * 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).
+     * Deletes by the full {@code (workspace_id, project_id, id)} sort key, matching {@code (project_id, id)} tuples so
+     * one statement can span several projects (OPIK-7483). Pairs are bound as two arrays and zipped, so the query text
+     * is constant regardless of batch size.
      * <p>
-     * {@code <if(partitions)>} adds the table's own weekly partition expression, bound as the set of partitions the
-     * batch's ids resolve to. It is emitted whenever every id in the batch is 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.
+     * {@code <if(partition)>} adds {@code IN PARTITION}, which scopes which <b>parts</b> the mutation is registered
+     * against — a {@code WHERE} clause cannot, since parts are selected before it is considered, which is why a delete
+     * of a few rows rewrote all ~3,650 parts and timed out. Omitted for the unbounded fallback.
      * <p>
-     * No schema flag gates it. {@link WeeklyPartitions} derives a value per {@code id_at} type the mutation may meet
-     * — the legacy 32-bit {@code DateTime} of {@code traces} as well as the {@code DateTime64(0)} of the partitioned
-     * successor — so one rendered statement is correct on both sides of the cutover EXCHANGE, in either direction, with
-     * nothing to flip and nothing to revert on rollback.
-     * <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
-     * as a literal. {@code arrayZip} is a deterministic function, not a subquery - ClickHouse rejects subqueries in
-     * delete mutations. Callers batch to keep each array within the driver's reliable bind size ({@link
-     * com.comet.opik.infrastructure.FilterUtils#ANALYTICS_DELETE_BATCH_SIZE}).
+     * {@code <partition>} is interpolated, not bound: {@code IN PARTITION {p:UInt32}} is a ClickHouse syntax error. Safe
+     * because the value is always a {@code long} from {@link WeeklyPartitions}.
      */
     private static final String DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS = """
             DELETE FROM <traces_mutation_table>
+            <if(partition)>IN PARTITION <partition><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>'
             ;
             """;
@@ -2328,7 +2310,7 @@ ORDER BY (workspace_id, project_id, id) DESC, last_updated_at DESC
      * {@code :min_id}'s week and inverted the window, so the fast pass matched nothing and every id fell through to
      * the unbounded pass. It now resolves them.
      * <p>
-     * The fallback remains load-bearing for what {@link com.comet.opik.utils.WeeklyPartitions#of} still cannot derive
+     * The fallback remains load-bearing for what {@link com.comet.opik.utils.WeeklyPartitions#groupByPartition} still cannot derive
      * exactly: an id at or past the end of {@code DateTime64}'s range, where {@code id_at} saturates to
      * {@code 2299-12-31} whatever the real week, so every such id collapses into one partition. Real data contains
      * them, so the bounded query is never a delete's sole resolver.
@@ -3645,42 +3627,79 @@ private Flux<? extends Result> getDetailsById(UUID id, Connection connection) {
     public Mono<Void> delete(Set<Pair<UUID, UUID>> projectIdTraceIdPairs, @NonNull Connection connection) {
         Preconditions.checkArgument(CollectionUtils.isNotEmpty(projectIdTraceIdPairs),
                 "Argument 'projectIdTraceIdPairs' must not be empty");
+        // Checked here rather than where the ids are stringified, so it holds whichever branch deleteBatch takes: the
+        // partitioned path would otherwise read a null as "underivable", silently take the unbounded fallback, and
+        // only then NPE - reporting a caller's bug as the slow delete this class exists to avoid.
+        Preconditions.checkArgument(
+                projectIdTraceIdPairs.stream().noneMatch(pair -> pair.getLeft() == null || pair.getRight() == null),
+                "Argument 'projectIdTraceIdPairs' must not contain null ids");
         log.info("Deleting traces by (project_id, id) pairs, count '{}'", projectIdTraceIdPairs.size());
 
         return makeMonoContextAware((userName, workspaceId) -> Flux
                 .fromIterable(Lists.partition(List.copyOf(projectIdTraceIdPairs), ANALYTICS_DELETE_BATCH_SIZE))
-                .concatMap(batch -> {
-                    var template = getSTWithLogComment(DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS, "delete_traces",
-                            workspaceId,
-                            userName, "pairs_size=%s".formatted(batch.size()));
-                    selectTracesMutationTable(template);
-
-                    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 every id in the batch allows it; otherwise emit the
-                    // unbounded form. Needs no schema flag: WeeklyPartitions derives a value per id_at type the
-                    // mutation may meet, so the set is correct on the legacy traces and on the partitioned successor.
-                    var partitions = WeeklyPartitions.of(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);
+                .concatMap(batch -> deleteBatch(batch, workspaceId, userName, connection))
+                .then());
+    }
 
-                    if (partitions.isPresent()) {
-                        statement = statement.bind("partitions", partitions.get().toArray(Long[]::new));
-                    }
+    /**
+     * Deletes one batch: one {@code IN PARTITION} statement per partition its ids resolve to, or a single unbounded
+     * statement when they cannot all be derived, or when the target is not partitioned.
+     * <p>
+     * Sequential on purpose: bounded concurrency was measured and deferred (OPIK-8230).
+     */
+    private Mono<Void> deleteBatch(List<Pair<UUID, UUID>> batch, String workspaceId, String userName,
+            Connection connection) {
+        // traceColumnsNonNullable doubles as "the mutation target is weekly-partitioned": the same cutover EXCHANGE
+        // drops the Nullable(...) columns and puts the partitioned successor behind the name mutations target, so one
+        // flag carries both facts. Deliberately not the wrap flag, which governs routing and is still false in the
+        // window between the EXCHANGE and the wrap - reading that one leaves production's deletes unpruned.
+        var grouped = traceColumnsNonNullable()
+                ? WeeklyPartitions.groupByPartition(batch.stream().map(Pair::getRight).toList())
+                : Optional.<Map<Long, Set<UUID>>>empty();
+
+        if (grouped.isEmpty()) {
+            return executeDelete(batch, null, workspaceId, userName, connection);
+        }
 
-                    var segment = startSegment("traces", "Clickhouse", "delete");
-                    return Mono.from(statement.execute())
-                            .doFinally(_ -> endSegment(segment))
-                            .then();
+        // id -> its pairs, built once per batch rather than rescanning the whole batch once per partition: an id can
+        // map to more than one pair (the same trace id reused across projects, OPIK-7483), so this is a
+        // Collectors.groupingBy, not a plain lookup map.
+        var pairsById = batch.stream().collect(Collectors.groupingBy(Pair::getRight));
+
+        return Flux.fromIterable(grouped.get().entrySet())
+                .concatMap(entry -> {
+                    var partitionPairs = entry.getValue().stream()
+                            .flatMap(id -> pairsById.get(id).stream())
+                            .toList();
+                    return executeDelete(partitionPairs, entry.getKey(), workspaceId, userName, connection);
                 })
-                .then());
+                .then();
+    }
+
+    /**
+     * Renders and executes one delete statement — unbounded when {@code partition} is null, scoped to it otherwise.
+     */
+    private Mono<Void> executeDelete(List<Pair<UUID, UUID>> pairs, Long partition,
+            String workspaceId, String userName, Connection connection) {
+        var template = getSTWithLogComment(DELETE_BY_PROJECT_ID_TRACE_ID_PAIRS, "delete_traces", workspaceId,
+                userName, "pairs_size=%s".formatted(pairs.size()));
+        selectTracesMutationTable(template);
+        if (partition != null) {
+            template.add("partition", partition);
+        }
+
+        var projectIds = pairs.stream().map(pair -> pair.getLeft().toString()).toArray(String[]::new);
+        var traceIds = pairs.stream().map(pair -> pair.getRight().toString()).toArray(String[]::new);
+
+        var statement = connection.createStatement(template.render())
+                .bind("workspace_id", workspaceId)
+                .bind("project_ids", projectIds)
+                .bind("trace_ids", traceIds);
+
+        var segment = startSegment("traces", "Clickhouse", "delete");
+        return Mono.from(statement.execute())
+                .doFinally(_ -> endSegment(segment))
+                .then();
     }
 
     /**

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

@@ -537,7 +537,7 @@ public Mono<Void> delete(@NonNull Set<UUID> ids, UUID projectId) {
      * Resolves every owning project for each id: a bounded fast pass, then an unbounded pass over only the ids the
      * bounded one leaves unresolved. Returns id -> owning projects; ids absent from the result have no live row.
      * <p>
-     * The bounded pass's week window can miss a row whose week {@link com.comet.opik.utils.WeeklyPartitions#of}
+     * The bounded pass's week window can miss a row whose week {@link com.comet.opik.utils.WeeklyPartitions#groupByPartition}
      * cannot derive exactly — an id at or past the end of {@code DateTime64}'s range, where {@code id_at} saturates
      * to {@code 2299-12-31} whatever the real week — so the unbounded pass re-resolves the miss set and the bounded
      * query is never a delete's sole resolver. A far-future timestamp short of that ceiling is no longer such a case:
@@ -566,6 +566,12 @@ private Mono<Map<UUID, Set<UUID>>> resolveOwningProjects(Set<UUID> ids) {
                 });
     }
 
+    /**
+     * All-or-nothing over the batch: an error anywhere skips {@code TracesDeleted} and the deletion-events capture for
+     * every pair, not just the failed one. OPIK-8230 widens the window — the DAO can now emit several statements per
+     * batch, so earlier partitions' rows may already be gone. Deletes are idempotent; restructuring this coupling is
+     * out of that ticket's scope.
+     */
     private Mono<Void> delete(Set<Pair<UUID, UUID>> projectIdTraceIdPairs, Connection connection) {
         return Mono.deferContextual(ctx -> {
             String workspaceId = ctx.get(RequestContext.WORKSPACE_ID);

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

@@ -15,6 +15,13 @@
  * epoch end time round-trips unchanged rather than being read as {@code null}. Flip this in lockstep with the EXCHANGE
  * step of the cutover.</p>
  *
+ * <p>It also gates partition-scoped deletes (OPIK-8230): the EXCHANGE that makes these columns non-nullable is the
+ * same one that puts the weekly-partitioned successor behind the name mutations target, so this flag being
+ * {@code true} is equally what says a delete may scope itself with {@code IN PARTITION}. One flag for two facts
+ * because they have only ever flipped together; a second would have to be threaded through the cutover runbook and
+ * tooling to track no independent state. The name says only the first duty and is deliberately not renamed - the env
+ * var is exposed.</p>
+ *
  * <p>{@code spanColumnsNonNullable}: the {@code spans} sibling of {@code traceColumnsNonNullable}, gating the same
  * sentinel wiring for {@code spans.end_time}→epoch and {@code spans.duration}/{@code spans.ttft}→{@code NaN}. Default
  * {@code false} while the {@code spans} table still has {@code Nullable(...)} columns; set {@code true} in lockstep with