# UNDF: UNDF-2026-000000716
# trino-0002: SkewedPartitionRebalancer scaledPartitions ArrayList.contains O(P²) per rebalance cycle

## Classification
- **Severity**: MEDIUM
- **CWE**: CWE-407 (Algorithmic Complexity — Inefficient Algorithmic Complexity)
- **Component**: `core/trino-main/src/main/java/io/trino/operator/output/SkewedPartitionRebalancer.java`
- **Method**: `rebalanceBasedOnTaskBucketSkewness()`

## Defect

`rebalanceBasedOnTaskBucketSkewness()` maintains `scaledPartitions` as `ArrayList<Integer>` to
track which partitions were already rebalanced in the current cycle. For each candidate partition
polled from a priority queue, it calls `scaledPartitions.contains(maxPartition)` — an O(P) linear
scan — where P is the number of scaled partitions so far.

The nested loop structure makes the total cost O(B × P²) per rebalance call, where B = number of
task buckets and P = number of partitions rebalanced this cycle. For wide tables or large fan-out
writes (P in the hundreds to thousands), this is quadratic per rebalance invocation.

`rebalance()` is called from `PartitionedOutputOperator.addInput()` whenever the output buffer is
full — i.e., continuously during a large-scale skewed write. This is a hot path in distributed
query execution.

### Defective code (lines 349, 376, 384)

```java
// line 349
List<Integer> scaledPartitions = new ArrayList<>();
while (true) {
    TaskBucket maxTaskBucket = maxTaskBuckets.poll();
    ...
    while (true) {
        Integer maxPartition = maxPartitions.poll();
        ...
        // line 376 — O(P) per iteration
        if (scaledPartitions.contains(maxPartition)) {
            continue;
        }
        ...
        scaledPartitions.add(maxPartition);  // line 384
    }
}
```

## Fix

Change `scaledPartitions` to `Set<Integer>` (use `new HashSet<>()`). The set semantics are
identical — it tracks which partitions have been scaled — but `contains()` becomes O(1).

```java
// Fix: O(1) dedup
Set<Integer> scaledPartitions = new HashSet<>();
...
if (scaledPartitions.contains(maxPartition)) {  // O(1)
    continue;
}
...
scaledPartitions.add(maxPartition);
```

## Complexity

| Partitions rebalanced (P) | Before (per cycle) | After (per cycle) |
|---------------------------|--------------------|--------------------|
| 100 | ~10,000 ops | ~100 ops |
| 1000 | ~1,000,000 ops | ~1,000 ops |
| 10000 | ~100,000,000 ops | ~10,000 ops |

**Speedup**: 100×–10,000× for large partition rebalance cycles.

## Patch

```diff
--- a/core/trino-main/src/main/java/io/trino/operator/output/SkewedPartitionRebalancer.java
+++ b/core/trino-main/src/main/java/io/trino/operator/output/SkewedPartitionRebalancer.java
@@ -1,2 +1,3 @@
+import java.util.HashSet;
+import java.util.Set;

 @GuardedBy("this")
 private void rebalanceBasedOnTaskBucketSkewness(...)
 {
-    List<Integer> scaledPartitions = new ArrayList<>();
+    Set<Integer> scaledPartitions = new HashSet<>();
     while (true) {
         ...
         while (true) {
             Integer maxPartition = maxPartitions.poll();
             ...
             if (scaledPartitions.contains(maxPartition)) {  // now O(1)
                 continue;
             }
```
