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

[VibhuJawa]

DRAFT for upstream rapidsai/cudf. Before filing: swap the VibhuJawa/cudf#9 cross-reference for its upstream number once that issue is filed, and delete this banner.

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<float32> 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<float64>	list<float32>	list<int8>
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<int8>[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<float32> 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<float32> → cuDF	NotImplementedError: large_list<item: float>	TypeError: Unsupported type: large_list<item: float>, 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

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(<cupy fp16 array>)	TypeError: Unsupported type float16
plc.Column.from_cuda_array_interface(<2-D fp16 array>)	ValueError: Unsupported dtype: f2
cudf.DataFrame.from_arrow on list<halffloat>	ValueError: … Unsupported type_id conversion to cudf (cpp/src/interop/arrow_utilities.cpp:61)
cudf.read_parquet of list<float16>	silently returns list<string>

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<float16> raises nothing and produces a list<string> 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:

"""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:

# 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 — [FEA] Make cudf::size_type 64-bit 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

• [FEA] The total number of characters in a cudf column is limited. Remove this restriction? rapidsai/cudf#12444 and [FEA] Increase maximum characters in strings columns 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.
• [FEA] Make cudf::size_type 64-bit rapidsai/cudf#3958 — [FEA] Make cudf::size_type 64-bit, closed 2022-06-25 with the wontfix label.
• [BUG] cudf.from_pandas breaks for large dataframes with list[..] column in 25.06+ rapidsai/cudf#18598 — cudf.from_pandas on large DataFrames with a list column; closed as fixed 2025-05-12.
• ParquetDatasetWriter max_file_size overestimates list columns and creates too many files 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.
• [FEA] No binary column type: Arrow binary/large_binary unsupported, Parquet binary silently becomes string #9 — cuDF has no binary column type. Same multimodal schema, separate root cause; tracked separately.

[VibhuJawa]

Filed upstream as rapidsai#23493