[tkislan]

@coderabbitai ignore

Summary

Fixes #141.

store_dataframe, added in 2.9.0, writes frames that can be loaded whole but not served a page at a time. Its manifest lists no per-batch row counts, batches are capped by rows only, and compression can fall back to lz4. With this PR the toolkit cuts the record batches itself and records each batch's row count in the manifest as batch_rows. Batches are capped by bytes as well as rows, lz4 is gone, and the manifest cap rises to 64 KiB.

Changes

• deepnote_toolkit/artifacts/_dataframe_storage.py:
  • One write path for pandas and polars.
    • pandas goes through pa.Table.from_pandas(df, preserve_index=None), the conversion write_feather did.
    • polars goes through df.to_arrow(), with compat_level=oldest from 1.2.1 as before.
    • The table is then written one batch at a time with pa.ipc.new_file, and the row count of each batch is returned.
  • Batch size. Rows per batch is clamp(8 MiB / average row bytes, 1, 65,536).
    • Row size is counted per column with nbytes, because every batch writes its own copy of each column and chunk, even ones that share buffers.
    • Columns holding a dictionary are the exception and use min(nbytes, get_total_buffer_size()). The file holds the dictionary once, while nbytes counts it once per chunk.
    • Each batch is a slice of the table with its chunks combined. Small chunks therefore don't make small batches, and the extra copy is at most one batch.
  • Dictionaries are unified before writing, because an IPC file can't replace a column's dictionary.
  • Compression is zstd, else none. Never lz4.
  • polars Object columns raise TypeError, including Object inside a List, Array or Struct. to_arrow() would export each value's memory address.
  • polars before 1.35.1: a frame with more than one chunk and a Categorical or Enum column is rechunked first. Those versions copy the whole dictionary into every chunk.
• deepnote_toolkit/artifacts/_dataframe_storage_manifest.py:
  • batch_rows is a new required field, and rows is its sum.
  • The cap goes from 4 KiB to 64 KiB.
  • encode_manifest raises ValueError past the cap, before anything is replaced.
• deepnote_toolkit/artifacts/dataframes.py: docstring.
• Tests in tests/unit/artifacts/.

Behaviour changes

• Manifests from 2.9.0: a manifest without batch_rows counts as absent. The next store of that name doesn't keep the old file as previous, and the old file is deleted once it is a day old, like any unreferenced frame. Frames stored by 2.9.0 can't be served a page at a time until they are stored again.
• Batch boundaries differ from 2.9.0. Values, schema and pandas metadata are unchanged. Where the boundaries match, the files are byte-identical (checked on pyarrow 10.0.1, 13, 16.1 and 24).
• Errors:
  • A frame that needs more than about 9,000 batches (some 600M rows) raises ValueError and keeps the previous frame. 2.9.0 never had this limit; without the guard, such a frame would get a manifest nothing can read.
  • polars Object columns raise TypeError instead of polars' ComputeError. On polars 1.39, write_ipc aborted the process for a List(Object) column; storing one now raises.
  • polars Int128 columns raise ArrowInvalid from to_arrow(). write_ipc used to write a file pyarrow couldn't read.
• polars memory: to_arrow() converts the whole frame at once, and copies string columns from views to large_string. On a 1.4 GiB frame of 10M strings the peak rose by 1.47 GiB, against 0.35 GiB with write_ipc. pandas is unchanged, because from_pandas always copied the whole frame.

Known limits

• Batch size comes from the frame's average row size, so very uneven rows can still make one batch several times 8 MiB.
• Categoricals stay dictionary-encoded, and a page reader has to read every dictionary for every page. Dictionaries past the reader's limit of about 1M values therefore make every page unservable. Left for a follow-up.
• A polars frame of about 100k chunks stores about 20 times slower than with write_ipc, because to_arrow() costs about 0.09 ms per chunk.
• The server-side reader must accept the 64 KiB manifest cap too.
• feat(dataframe-storage): store a block's result automatically #131 changes the same files and will need a rebase.

Testing

• Storage tests: 127 pass on Python 3.13 (pyarrow 24.0.0, pandas 2.2.3, polars 1.39.3) and on Python 3.10 (pyarrow 16.1.0, pandas 2.1.4).

• Full unit suite: 1,406 passed and 10 skipped, with no new warnings.

• Mutation check: each of 24 mutants of the new code fails at least one test. They include:

  • a missing chunk merge, unify, byte cap or slice clamp;
  • a 125-row floor;
  • an lz4 fallback;
  • a wrong row-size rule for shared buffers or dictionaries;
  • a missing Object check or Struct recursion;
  • the rechunk gate off;
  • a missing manifest guard.

• Version checks in throwaway environments: pyarrow 10.0.1, 13.0.0, 16.1.0 and 24.0.0 with pandas 1.5.3 to 2.2.3, and polars 0.20.31 to 1.39.3.

• End to end: a server-side page reader, with its manifest cap raised to 64 KiB, served every page at pageSize 1000 and 100 for frames written by this branch. The frames covered:

  • ints and categoricals;
  • polars frames with small chunks, and with 4,000 chunks of a categorical (with a whole-table nbytes estimate, 0 of 400 pages were served);
  • frames whose columns or chunks share buffers: pd.concat([df] * 5), pl.concat([df] * 10) and aliased columns;
  • empty frames and frames without columns;
  • long text;
  • uncompressed incompressible bytes, which fail at a 16 MiB batch target;
  • mixed types.

  Wide frames hit the reader's per-page value limit only at pageSize 1000.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VNWoySEbnsryRwTGKUydNx

[github-actions]

📦 Python package built successfully!

• Version: 2.9.0.dev2+f9bf9d3
• Wheel: deepnote_toolkit-2.9.0.dev2+f9bf9d3-py3-none-any.whl
• Install:

  pip install "deepnote-toolkit @ https://deepnote-staging-runtime-artifactory.s3.amazonaws.com/deepnote-toolkit-packages/2.9.0.dev2%2Bf9bf9d3/deepnote_toolkit-2.9.0.dev2%2Bf9bf9d3-py3-none-any.whl"

[codecov]

Codecov Report

❌ Patch coverage is 97.14286% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.30%. Comparing base (11a8bed) to head (f1c6735).
✅ All tests successful. No failed tests found.

Files with missing lines	Patch %	Lines
deepnote_toolkit/artifacts/_dataframe_storage.py	96.72%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #142      +/-   ##
==========================================
+ Coverage   78.17%   78.30%   +0.12%     
==========================================
  Files         120      120              
  Lines        6837     6876      +39     
  Branches      990      997       +7     
==========================================
+ Hits         5345     5384      +39     
  Misses       1182     1182              
  Partials      310      310              

Flag	Coverage Δ
combined	78.30% <97.14%> (+0.12%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.