NVIDIA / NVIDIA/cudf

[FEA] Support >2^31 elements in list column children (large-offset list columns)

Open
#23,493 2 comments 0 reactions 0 assignees View on GitHub
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

## Is your feature request related to a problem? Please describe.

`cudf::size_type` is `int32_t`, so a column is capped at 2^31−1 ≈ 2.1 B elements. For a `LIST` column that cap applies to the **leaf/child column**, not to the row count. A `list` with `D` values per row therefore tops out at `floor((2^31−1) / D)` rows.

Crucially, the cap counts **elements**, not bytes, so the row ceiling is independent of the leaf dtype:

```
max_rows(D) = floor((2^31 - 1) / D) # no sizeof(T) anywhere in this
leaf_payload = max_rows * D * sizeof(T)
offsets = (max_rows + 1) * 4 bytes # int32 offsets child
```

| dim `D` | max rows per frame | leaf elements at cap | device mem `list` | `list` | `list` |
|---:|---:|---:|---:|---:|---:|
| 128 | 16,777,215 | 2,147,483,520 | 16.06 GiB | 8.06 GiB | 2.06 GiB |
| 256 | 8,388,607 | 2,147,483,392 | 16.03 GiB | 8.03 GiB | 2.03 GiB |
| 384 | 5,592,405 | 2,147,483,520 | 16.02 GiB | 8.02 GiB | 2.02 GiB |
| 512 | 4,194,303 | 2,147,483,136 | 16.02 GiB | 8.02 GiB | 2.02 GiB |
| 768 | 2,796,202 | 2,147,483,136 | 16.01 GiB | 8.01 GiB | 2.01 GiB |
| 1024 | **2,097,151** | 2,147,482,624 | 16.01 GiB | 8.01 GiB | 2.01 GiB |
| 1536 | 1,398,101 | 2,147,483,136 | 16.01 GiB | 8.01 GiB | 2.01 GiB |
| 2048 | 1,048,575 | 2,147,481,600 | 16.00 GiB | 8.00 GiB | 2.00 GiB |
| 4096 | 524,287 | 2,147,479,552 | 16.00 GiB | 8.00 GiB | 2.00 GiB |

Read across any row: the row count never moves. Halving the element width halves the memory and buys zero additional rows. An `int8` embedding column — the quantized format reached for *precisely* to fit more rows on a GPU — stops at the same row count as `float64` while holding an eighth of the bytes. On an 80 GB device a maxed-out `list[1024]` column occupies **2 GiB of 80**, and the next row still fails. Unlike the other memory ceilings in the stack, this one does not respond to dtype choice.

An embedding column is the most common `list` in AI data pipelines, and ~2 M rows per frame is small — we routinely process hundreds of millions of embedding rows per job. This is the constraint we hit first and most often.

### Verified behaviour

Environment: NVIDIA H100 80GB HBM3 (79.2 GiB total, **78.7 GiB free**), cupy 14.1.1, `dim=1024`, each case allocating ~8.00 GiB. Run on **cuDF 25.10.00** (pyarrow 25.0.0) and **26.08.00a990 nightly** (pyarrow 23.0.1).

Each probe has a just-under-cap **control** and a just-over-cap case. Controls pass and over-cap cases fail with ~70 GiB still free, so these are attributable to the cap, not to memory pressure.

| # | Probe | 25.10.00 | 26.08.00a990 |
|---|---|---|---|
| 1a | Build `LIST`, leaf **2,147,482,624** (just under) | OK | OK |
| 1b | Build `LIST`, leaf **2,147,483,648** (just over) | `ValueError: Flat size exceeds size_type limit` | same |
| 2a | `cudf.concat` to leaf just under | OK | OK |
| 2b | `cudf.concat` to leaf just over | `OverflowError: Total number of concatenated rows exceeds the column size limit` (`cpp/src/copying/concatenate.cu:478`) | same, same line |
| 3 | Arrow `large_list` → cuDF | `NotImplementedError: large_list` | `TypeError: Unsupported type: large_list`, chained from the same `NotImplementedError` |

Behaviour is the same on both versions. Two cosmetic differences: probe 3's exception is re-wrapped as `TypeError` in 26.08, and the Python-side traceback line numbers moved (`pylibcudf/column.pyx:972`/`:302` in 25.10 → `:1059`/`:307` in 26.08).

**The element cap is dtype-independent — measured, not just derived.** `dim=1024`, column construction only, no I/O, cuDF 26.08.00a990:

| dtype | rows | leaf elements | payload | result |
|---|---:|---:|---:|---|
| `float64` | 2,097,151 | 2,147,482,624 | 16.00 GiB | builds |
| `float64` | 2,097,152 | 2,147,483,648 | 16.00 GiB | `ValueError: Flat size exceeds size_type limit` |
| `float32` | 2,097,151 | 2,147,482,624 | 8.00 GiB | builds |
| `float32` | 2,097,152 | 2,147,483,648 | 8.00 GiB | same `ValueError` |
| `int8` | 2,097,151 | 2,147,482,624 | 2.00 GiB | builds |
| `int8` | 2,097,152 | 2,147,483,648 | 2.00 GiB | same `ValueError` |

All three fail at exactly the same row — 2,097,152 — across an 8× range in bytes actually held.

Tracebacks (1b and 2b, cuDF 26.08.00a990)

```
File "repro_list_cap.py", line 44, in make_list_series
return cudf.Series.from_pylibcudf(plc.Column.from_cuda_array_interface(arr))
File "pylibcudf/column.pyx", line 1059, in pylibcudf.column.Column.from_cuda_array_interface
File "pylibcudf/column.pyx", line 307, in pylibcudf.column._prepare_array_metadata
ValueError: Flat size exceeds size_type limit
```

```
File "cudf/core/dataframe.py", line 2235, in _concat
plc_result = plc.concatenate.concatenate(plc_tables)
File "pylibcudf/concatenate.pyx", line 57, in pylibcudf.concatenate.concatenate
OverflowError: CUDF failure at: /__w/cudf/cudf/cpp/src/copying/concatenate.cu:478:
Total number of concatenated rows exceeds the column size limit
```

Steps/code to reproduce

```python
import traceback
import cupy as cp
import cudf
import pylibcudf as plc

SIZE_TYPE_MAX = 2**31 - 1
DIM = 1024
ROWS_UNDER = SIZE_TYPE_MAX // DIM # 2,097,151 -> leaf 2,147,482,624
ROWS_OVER = ROWS_UNDER + 1 # 2,097,152 -> leaf 2,147,483,648

def probe(name, fn):
try:
print(f"{name}: OK ->", fn())
except Exception as exc:
print(f"{name}: RAISED {type(exc).__name__}")
traceback.print_exc()
cp.get_default_memory_pool().free_all_blocks()

def make_list_series(rows, dtype="float32"):
arr = cp.zeros((rows, DIM), dtype=dtype)
return cudf.Series.from_pylibcudf(plc.Column.from_cuda_array_interface(arr))

def concat_case(total_rows):
half = total_rows // 2
a = cudf.DataFrame({"e": make_list_series(half)})
b = cudf.DataFrame({"e": make_list_series(total_rows - half)})
return len(cudf.concat([a, b], ignore_index=True))

probe("1a control", lambda: len(make_list_series(ROWS_UNDER)))
probe("1b over ", lambda: len(make_list_series(ROWS_OVER)))
probe("2a control", lambda: concat_case(ROWS_UNDER))
probe("2b over ", lambda: concat_case(ROWS_OVER))

# dtype independence: all three fail at the same row count
for dt in ("float64", "float32", "int8"):
probe(f"{dt} at cap ", lambda dt=dt: len(make_list_series(ROWS_UNDER, dt)))
probe(f"{dt} cap + 1", lambda dt=dt: len(make_list_series(ROWS_OVER, dt)))
```

### There is no narrower dtype to fall back to

`float16` and `bfloat16` are not supported by cuDF, so the only rung below `float32` is `int8` (verified on 26.08.00a990):

| construction path | `float16` result |
|---|---|
| `cudf.Series()` | `TypeError: Unsupported type float16` |
| `plc.Column.from_cuda_array_interface(<2-D fp16 array>)` | `ValueError: Unsupported dtype: f2` |
| `cudf.DataFrame.from_arrow` on `list` | `ValueError: … Unsupported type_id conversion to cudf` (`cpp/src/interop/arrow_utilities.cpp:61`) |
| `cudf.read_parquet` of `list` | **silently returns `list`** |

`cudf.utils.dtypes.SUPPORTED_NUMPY_TO_PYLIBCUDF_TYPES` contains no `float16` and no `bfloat16` at any width.

That last row is a separate small trap: reading a Parquet `list` raises nothing and produces a `list` column whose "strings" are the raw 2-byte half-float values (`.str.byte_count()` is `2` for every element; reinterpreting those bytes as `float16` recovers the original values bit-exactly). This is consistent with cuDF surfacing the `FIXED_LEN_BYTE_ARRAY(2)` physical type and ignoring the `FLOAT16` logical annotation. A flat, non-list `float16` Parquet column likewise comes back with `list` dtype rather than anything numeric.

So the practical dtype ladder for an embedding column is `float64` → `float32` → `int8`, and no rung on it changes the row ceiling.

### What we do downstream today

NeMo Curator carries two workarounds.

**1.** A file-regrouping pass whose only purpose is to keep each cuDF frame under the cap — [`break_parquet_partition_into_groups`](https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/deduplication/semantic/utils.py):

```python
"""Break parquet files into groups to avoid cudf 2bn row limit."""
if embedding_dim is None:
embedding_dim = 1024 # default aggressive assumption
cudf_max_num_rows = 2_000_000_000 # cudf only allows 2bn rows
cudf_max_num_elements = cudf_max_num_rows / embedding_dim
# cudf considers each element in an array to be a row
```

It reads only the **first** file's Parquet footer and extrapolates row counts with a hard-coded **1.5× skew factor**, because there is no way to ask cuDF "will this fit?" ahead of time. Sufficiently skewed inputs still overflow.

**2.** Concatenation is split apart so embeddings never pass through `cudf.concat` — [`pairwise.py`](https://github.com/NVIDIA-NeMo/Curator/blob/main/nemo_curator/stages/deduplication/semantic/pairwise.py):

```python
# Cannot concatenate dataframes with embeddings due to cudf 2bn row limit
# Instead, concatenate metadata columns and handle embeddings separately
```

Probe 2b above is that comment, reproduced.

Net effect: embeddings are moved out of cuDF into CuPy/torch as early as possible, and cuDF is used only for metadata columns plus Parquet I/O. We lose cuDF for the widest column in the table.

## Describe the solution you'd like

**Large-offset `LIST` columns**, solved the way large strings were: promote the offsets child to `int64` when the leaf exceeds `size_type` range, while the row count stays in `size_type`. It would need to hold through `lists_column_view`, Parquet read **and write**, `concat` / `gather` / `explode` / `.list.leaves`, and Arrow `large_list` interop.

The row-count cap is fine for us. It is specifically the leaf-element cap that binds.

## Describe alternatives you've considered

- **Making `cudf::size_type` 64-bit globally** — rapidsai/cudf#3958, closed with the `wontfix` label. Explicitly *not* what this asks for.
- **Chunking above cuDF** — what we do today. Fragile, and it caps per-frame parallelism well below what the GPU can hold.
- **Keeping embeddings in CuPy/torch** — also what we do today. The embedding column never benefits from cuDF ops, and it forces a split between metadata (cuDF) and payload (CuPy).
- **A narrower leaf dtype** — buys nothing; see the tables above.

## Additional context

- rapidsai/cudf#12444 and rapidsai/cudf#13733 — the string character cap, both closed as completed and resolved by moving strings to 64-bit offsets. **This is the direct analogue for lists.**
- rapidsai/cudf#3958 — `[FEA] Make cudf::size_type 64-bit`, closed 2022-06-25 with the `wontfix` label.
- rapidsai/cudf#18598 — `cudf.from_pandas` on large DataFrames with a list column; closed as fixed 2025-05-12.
- rapidsai/cudf#23378 — `ParquetDatasetWriter max_file_size` overestimates list columns and creates too many files (open); another place the list leaf/row distinction is mishandled.
- #23492 — cuDF has no binary column type. Same multimodal schema, separate root cause; tracked separately.

Contributor guide

Open the contributing guide

Research direction

Start by tracing lists_column_view and the concatenation failure at cpp/src/copying/concatenate.cu:478, then compare the existing large-string 64-bit offset handling described in the issue. Map the affected paths through Parquet I/O, concat, gather, explode, .list.leaves, and Arrow large_list interop. Done means large-offset LIST columns work across those paths while the row count remains size_type-limited.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
data
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.