feat(python): support sequence packing in streaming dataset (#3920)

## How packing works

  Consider four tokenized documents:

  
  [1]
  [2]
  [10, 11, 12, 13, 14, 15, 16, 17]
  [20]

  With:
```
  StreamingDataset(
      table,
      shuffle=False,
      columns=["tokens"],
      num_splits=2,
      pack_sequences=5,
      eos_id=9,
      pad_id=0,
      blocks_per_epoch=6,
  )
```
the documents are assigned to two fixed logical splits. Each split
maintains an independent token buffer, appends eos_id after every
document, and emits blocks of five tokens.

  Because blocks_per_epoch=6, each split emits exactly three blocks:

  Cycl/e 1:
    Split 0: [1, 9, 2, 9, 0] # 9 is eos, 0 is padding
    Split 1: [10, 11, 12, 13, 14]

  Cycle 2:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [15, 16, 17, 9, 20]

  Cycle 3:
    Split 0: [0, 0, 0, 0, 0]
    Split 1: [9, 0, 0, 0, 0]

If a split runs out of tokens early, it emits padded blocks through the
fixed budget. This prevents one rank from finishing before another.

Logical splits are independent of rank and worker ownership. A
checkpoint records each split’s consumed-document count, emitted-block
count, remaining tokens, and document boundaries. Merging
those per-split states allows the same packed stream to resume after the
topology changes.

doc_ids identifies document segments, including continuations across
block boundaries. It is not a padding mask: padding retains the
preceding document ID, so callers must mask padding using a
  reserved pad_id.

blocks_per_epoch="auto" is also available. It estimates the budget from
a deterministic bounded sample and warns that the result is approximate.

WIP pre-training tests:
```
  ┌────────────────────────────────────┬────────────────────────┬─────────────────────────────────────┐
  │                                    │       GPT-2 124M       │          GPT-2 medium 354M          │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Corpus                             │ 2.4M docs / 12GB table │ 9.67M docs / 45GB table             │
  │ Tokens (Chinchilla)                │ 2.43B                  │ 7.0B                                │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Data prep (ingest→curate→tokenize) │ ~12 min                │ ~51 min                             │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Training wall time                 │ ~50 min                │ 3h 06m                              │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Throughput / MFU                   │ 1.60M tok/s / 35%      │ 684k tok/s / 42.0%, │
  ├────────────────────────────────────┼────────────────────────┼─────────────────────────────────────┤
  │ Final val loss                     │ 3.230                  │ 2.840                               │
  └────────────────────────────────────┴────────────────────────┴─────────────────────────────────────┘
```

---------

Co-authored-by: OpenAI Codex <codex@openai.com>

python/python/tests/test_elastic_dataloader.py

@@ -2113,6 +2113,214 @@ def test_shuffle_seed_none_generates_stable_seed(lance_table):
     assert first == second, "Same resolved seed must produce the same ordering"
 
 
+# Sequence packing tests
+
+
+def _create_token_table(tmp_path, documents):
+    db = lancedb.connect(tmp_path)
+    tokens = pa.array(documents, type=pa.list_(pa.int64()))
+    return db.create_table("tokens", pa.table({"tokens": tokens}))
+
+
+def _packed_dataset(table, pack_sequences, *, blocks_per_epoch, pad_id=0, **kwargs):
+    return StreamingDataset(
+        table,
+        shuffle=False,
+        columns=["tokens"],
+        pack_sequences=pack_sequences,
+        eos_id=9,
+        pad_id=pad_id,
+        blocks_per_epoch=blocks_per_epoch,
+        **kwargs,
+    )
+
+
+def test_pack_sequences_emits_blocks_and_pads_final_tail(tmp_path):
+    table = _create_token_table(tmp_path, [[1, 2], [3, 4], [5]])
+    dataset = _packed_dataset(table, 6, blocks_per_epoch=2)
+
+    blocks = list(dataset)
+
+    assert len(blocks) == 2
+    assert blocks[0]["input_ids"].tolist() == [1, 2, 9, 3, 4, 9]
+    assert blocks[0]["doc_ids"].tolist() == [0, 0, 0, 1, 1, 1]
+    assert blocks[1]["input_ids"].tolist() == [5, 9, 0, 0, 0, 0]
+    assert blocks[1]["doc_ids"].tolist() == [0, 0, 0, 0, 0, 0]
+    assert blocks[0]["input_ids"].dtype == torch.int64
+    assert blocks[0]["doc_ids"].dtype == torch.int64
+
+
+def test_pack_sequences_pads_lagging_splits(tmp_path):
+    table = _create_token_table(
+        tmp_path,
+        [[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
+    )
+    dataset = _packed_dataset(table, 5, blocks_per_epoch=6, num_splits=2)
+    input_ids = [block["input_ids"].tolist() for block in dataset]
+    # Split 0 has four real tokens including EOS markers, while split 1 has
+    # eleven. Packing must emit three complete two-split cycles.
+    assert input_ids == [
+        [1, 9, 2, 9, 0],
+        [10, 11, 12, 13, 14],
+        [0, 0, 0, 0, 0],
+        [15, 16, 17, 9, 20],
+        [0, 0, 0, 0, 0],
+        [9, 0, 0, 0, 0],
+    ]
+
+    per_rank = []
+    for rank in range(2):
+        rank_dataset = _packed_dataset(
+            table,
+            5,
+            blocks_per_epoch=6,
+            num_splits=2,
+            world_size=2,
+            rank=rank,
+        )
+        per_rank.append([block["input_ids"].tolist() for block in rank_dataset])
+
+    assert [len(blocks) for blocks in per_rank] == [3, 3]
+    sharded = [block for cycle in zip(*per_rank) for block in cycle]
+    assert sharded == input_ids
+
+
+def test_pack_sequences_auto_estimates_filtered_token_column(tmp_path):
+    db = lancedb.connect(tmp_path)
+    table = db.create_table(
+        "tokens",
+        pa.table(
+            {
+                "tokens": pa.array([[1] * 4, [2] * 9], type=pa.list_(pa.int64())),
+                "keep": [True, False],
+            }
+        ),
+    )
+    table.add(
+        pa.table(
+            {
+                "tokens": pa.array([[3] * 4, [4] * 9], type=pa.list_(pa.int64())),
+                "keep": [True, False],
+            }
+        )
+    )
+
+    with pytest.warns(UserWarning, match="approximate token-count sample"):
+        dataset = _packed_dataset(
+            table,
+            5,
+            blocks_per_epoch="auto",
+            num_splits=2,
+            filter="keep",
+        )
+
+    # Two kept documents contain 8 tokens plus 2 EOS tokens: two blocks.
+    assert dataset.state_dict()["blocks_per_epoch"] == 2
+
+
+def test_pack_sequences_checkpoint_resumes_on_new_topology(tmp_path):
+    table = _create_token_table(
+        tmp_path,
+        [[1], [2], [10, 11, 12, 13, 14, 15, 16, 17], [20]],
+    )
+    kwargs = dict(pack_sequences=5, blocks_per_epoch=6, num_splits=2)
+    reference = list(_packed_dataset(table, **kwargs))
+
+    datasets = [
+        _packed_dataset(table, world_size=2, rank=rank, **kwargs) for rank in range(2)
+    ]
+    iterators = [iter(dataset) for dataset in datasets]
+    first_cycle = [next(iterator) for iterator in iterators]
+    checkpoint = StreamingDataset.merge_state_dicts(
+        [dataset.state_dict() for dataset in datasets]
+    )
+    for iterator in iterators:
+        iterator.close()
+
+    resumed = _packed_dataset(table, **kwargs)
+    resumed.load_state_dict(checkpoint)
+    actual_remaining = list(resumed)
+
+    assert [block["input_ids"].tolist() for block in first_cycle] == [
+        [1, 9, 2, 9, 0],
+        [10, 11, 12, 13, 14],
+    ]
+    assert checkpoint["blocks_emitted_per_split"] == [1, 1]
+    assert [block["input_ids"].tolist() for block in actual_remaining] == [
+        block["input_ids"].tolist() for block in reference[2:]
+    ]
+    assert [block["doc_ids"].tolist() for block in actual_remaining] == [
+        block["doc_ids"].tolist() for block in reference[2:]
+    ]
+
+
+def test_pack_sequences_validates_configuration_and_tokens(tmp_path):
+    table = _create_token_table(tmp_path, [[1, 2]])
+
+    with pytest.raises(ValueError, match="pad_id is required"):
+        StreamingDataset(
+            table,
+            shuffle=False,
+            columns=["tokens"],
+            pack_sequences=4,
+            eos_id=9,
+        )
+
+    with pytest.raises(ValueError, match="blocks_per_epoch is required"):
+        StreamingDataset(
+            table,
+            shuffle=False,
+            columns=["tokens"],
+            pack_sequences=4,
+            eos_id=9,
+            pad_id=0,
+        )
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        _packed_dataset(table, 4, blocks_per_epoch=3, num_splits=2)
+
+    with pytest.raises(ValueError, match="positive integer or 'auto'"):
+        _packed_dataset(table, 4, blocks_per_epoch="estimate")
+
+    checkpoint = _packed_dataset(table, 4, blocks_per_epoch=1).state_dict()
+    resumed = _packed_dataset(table, 4, blocks_per_epoch=1, pad_id=8)
+    with pytest.raises(ValueError, match="pad_id mismatch"):
+        resumed.load_state_dict(checkpoint)
+
+    float_db = lancedb.connect(tmp_path / "float")
+    float_table = float_db.create_table(
+        "tokens",
+        pa.table({"tokens": pa.array([[1.5, 2.5]], type=pa.list_(pa.float64()))}),
+    )
+    with pytest.raises(ValueError, match="token column with integer values"):
+        _packed_dataset(float_table, 4, blocks_per_epoch=1)
+
+    null_db = lancedb.connect(tmp_path / "null")
+    null_table = null_db.create_table(
+        "tokens",
+        pa.table({"tokens": pa.array([None], type=pa.list_(pa.int64()))}),
+    )
+    with pytest.raises(ValueError, match="does not support null token lists"):
+        list(_packed_dataset(null_table, 4, blocks_per_epoch=1))
+
+    null_value_db = lancedb.connect(tmp_path / "null_value")
+    null_value_table = null_value_db.create_table(
+        "tokens",
+        pa.table(
+            {"tokens": pa.array([[1], [2, None], [3]], type=pa.list_(pa.int64()))}
+        ),
+    )
+    blocks = list(
+        _packed_dataset(
+            null_value_table,
+            2,
+            blocks_per_epoch=2,
+            on_transform_error="skip",
+        )
+    )
+    assert [block["input_ids"].tolist() for block in blocks] == [[1, 9], [3, 9]]
+
+
 # ---------------------------------------------------------------------------
 # Doc examples — each test mirrors the code snippet in index.mdx so that
 # broken doc examples are caught before they ship.