kubeflow / kubeflow/trainer

[BUG]: IndexableMemTable underflows on an empty batch and fetch_partitions can return a non-overlapping batch

Open
#4,056 1 comment 0 reactions 1 assignee Claimed by @reckless-sherixx View on GitHub
Dominant language
Go
Stars
2.2k
Forks
1.1k
Avg merge
3d 22h
Merged PRs (30d)
39

Description

### What happened?

Two index-arithmetic defects in `IndexableMemTable` (`pkg/data_cache/src/worker/indexable_mem_table.rs`).

#### (a) `load` underflows on an empty batch

```rust
// indexable_mem_table.rs:82
let current_end = current_start + num_rows - 1;
```

With `num_rows == 0` and `current_start == 0` this evaluates `0 + 0 - 1`. With overflow checks on it panics. The shipped image is built `--release` (`cmd/data_cache/Dockerfile:42`), where it wraps to `u64::MAX`; `current_start` then wraps back to 0 and `end_indices` is left unsorted. Every later `fetch_partitions` call binary searches `start_indices` and `end_indices` with `partition_point`, so from that point on the worker serves arbitrary row ranges — silently, with no error anywhere.

Running that arithmetic over a scan that yields batches of `[0, 4, 4]` rows starting at index 0:

```
debug (overflow checks on) -> panicked at 'attempt to subtract with overflow'
release (overflow checks off) -> start_indices = [0, 0, 4]
end_indices = [18446744073709551615, 3, 7]
end_indices sorted = false
```

Empty batches do reach this loop: `RowNumberStream` passes them straight through (`pkg/data_cache/src/worker/worker_datasource.rs:481-499`), and Iceberg delete-file filtering can empty a batch.

#### (b) `fetch_partitions` can return a batch that does not overlap the query

```rust
// indexable_mem_table.rs:120-128
let last = start_indices
.partition_point(|&batch_start| batch_start <= end)
.saturating_sub(1);

// Verify we have valid range
if first > last {
return vec![];
}
```

When no batch starts at or below `end`, `partition_point` returns 0 and `saturating_sub(1)` clamps it to 0 — indistinguishable from "batch 0 qualifies". The `first > last` guard then passes and batch 0 is returned even though it lies entirely outside the requested range:

```
batches covering global rows 10-13 and 14-17, query [0, 5] -> 1 batch returned
```

The impact today is bounded: `supports_filters_pushdown` returns `Inexact` (`indexable_mem_table.rs:199`), so DataFusion keeps the `FilterExec` and discards the extra rows. It is wasted work now and a correctness bug the moment pushdown is tightened to `Exact`.

### What did you expect to happen?

(a) An empty batch owns no rows, so it should be skipped rather than given an index range. The index vectors have to stay sorted for the binary searches that read them.

(b) When no batch starts at or below `end`, `fetch_partitions` should return `vec![]`.

### Existing test coverage

None for either case. `load`'s tests feed a `MemTable` with two four-row batches (`indexable_mem_table.rs:368`), and all four `fetch_partitions` tests (`:274`, `:295`, `:316`, `:399`) use query ranges that overlap the batches.

### Environment

Reproduced from source at `master` @ `8ca43ec` (VERSION `v2.3.0`) by running `cargo test` against `pkg/data_cache`. Case (a) panics under `cargo test`; the release figures above come from compiling the same arithmetic with `-C overflow-checks=off`, matching the profile the image ships.

Kubernetes version:
```bash
$ kubectl version
n/a - reproduced from source
```

Kubeflow Trainer version:
```bash
$ kubectl get pods -n kubeflow-system -l app.kubernetes.io/name=kubeflow-trainer -o jsonpath="{.items[*].spec.containers[*].image}"
n/a - reproduced from source at master @ 8ca43ec (VERSION v2.3.0)
```

Kubeflow Python SDK version:
```bash
$ pip show kubeflow
n/a
```

### Impacted by this bug?

Give it a 👍 We prioritize the issues with most 👍

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.