From 5111419832c0a6e43857b1bbe412a977389c0ea9 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Wed, 4 Nov 2026 06:48:32 +0900 Subject: [PATCH 1/2] Native iter_json_batches reads the cursor one batch at a time (closes #239) iter_json_batches drained the whole result into dicts (_drain_dicts), converted every row, and only then cut the list into batches, so batch_size bounded nothing and the peak held every row twice. It now holds at most one reply (_NEXT_BATCH rows) and one output batch, and reads the cursor when the next batch is asked for. Rows, order, types and batch boundaries are those of to_json_list cut at batch_size. A result closed between two batches raises the closed-before-read error, as the other iterators do. to_columns, to_arrow and to_polars are unchanged (a column type is decided over all rows). Co-Authored-By: Claude Opus 5.5 --- src/arcadedb/_native/database.py | 35 ++++++- tests/conftest.py | 2 + tests/test_native_streaming_batches.py | 134 +++++++++++++++++++++++++ 3 files changed, 168 insertions(+), 3 deletions(-) create mode 100644 tests/test_native_streaming_batches.py diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index 8fb321f..865dd11 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -283,6 +283,36 @@ def _drain_dicts(self) -> List[Dict[str, Any]]: self._finish() return out + def _iter_dict_batches(self, size: int) -> Iterator[List[Dict[str, Any]]]: + """The rows as lists of ``size`` dicts (the last may be shorter), reading the cursor one batch at a time: at most one reply + (``_NEXT_BATCH`` rows) and one output batch are held, where ``_drain_dicts`` holds the whole result. The cursor is read + when the next batch is asked for, so a result closed between two batches raises the closed-before-read error. + """ + if not self._readable(): + return + out: List[Dict[str, Any]] = [] + while True: + buf = self._buf + while buf: + need = size - len(out) + if len(buf) <= need: + out.extend([p for p, _ in buf]) + buf.clear() + else: + out.extend([buf.popleft()[0] for _ in range(need)]) + if len(out) >= size: + full, out = out, [] + if not buf and not self._more and self._pending_error is None: + self._finish() # this batch is the last row: the cursor is finished before it is handed over + yield full + if not self._readable(): + return + if not self._fill(): + break + self._finish() + if out: + yield out + def to_list(self, convert_types: bool = True) -> List[Dict[str, Any]]: return self._drain_dicts() @@ -298,9 +328,8 @@ def to_json_list(self, batch_size: int = 10_000) -> List[Dict[str, Any]]: def iter_json_batches(self, batch_size: int = 10_000, **kwargs): if int(batch_size) < 1: raise ValueError("batch_size must be at least 1") - rows = self.to_json_list() - for i in range(0, len(rows), batch_size): - yield rows[i : i + batch_size] + for batch in self._iter_dict_batches(int(batch_size)): + yield [json_native(row) for row in batch] def to_dataframe(self, convert_types: bool = True): import pandas as pd diff --git a/tests/conftest.py b/tests/conftest.py index 01ab332..5b7be31 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -63,6 +63,8 @@ collect_ignore.append("test_native_tx_crossings.py") # counts the statement texts the native backend hands to the engine (the JPype backend has no such layer) collect_ignore.append("test_native_statement_texts.py") + # iter_json_batches reads the cursor one batch at a time (the result set's buffer and its _fill are observed) + collect_ignore.append("test_native_streaming_batches.py") if not os.path.isfile( os.path.join(os.path.dirname(os.path.dirname(__file__)), "docs", "index.md") ): diff --git a/tests/test_native_streaming_batches.py b/tests/test_native_streaming_batches.py new file mode 100644 index 0000000..e864bfb --- /dev/null +++ b/tests/test_native_streaming_batches.py @@ -0,0 +1,134 @@ +"""``iter_json_batches`` on the native build reads the cursor one batch at a time (#239). + +It used to drain the whole result into dicts, convert every row, and only then cut the list into batches, so ``batch_size`` bounded nothing: +the peak held every row of the result. Now at most one reply of the cursor (``_NEXT_BATCH`` rows) and one output batch are held, and the +cursor is read when the next batch is asked for. The rows, their order, the batch boundaries and the types are those of ``to_json_list`` +(compared here on a result of several replies: ``arcadedb_n_exec`` returns the first 256 rows, ``arcadedb_n_next`` the following 2048 at a +time). conftest.py collects this file on the native backend only: it observes the result set's buffer, which the JVM build does not have. +""" + +import pytest + +import arcadedb +from arcadedb._native import database + +ROWS = 5000 # 256 in the first reply, then 2048 + 2048 + 648 + + +@pytest.fixture +def db(temp_db): + temp_db.command("sql", "CREATE DOCUMENT TYPE R") + temp_db.command("sql", "CREATE PROPERTY R.id INTEGER") + temp_db.command("sql", "CREATE PROPERTY R.txt STRING") + temp_db.command("sql", "CREATE PROPERTY R.val DOUBLE") + temp_db.insert_many( + "R", [{"id": i, "txt": f"t{i}", "val": i / 8.0} for i in range(ROWS)] + ) + return temp_db + + +def _select(db): + return db.query("sql", "SELECT id, txt, val FROM R ORDER BY id") + + +@pytest.fixture +def fills(monkeypatch): + """The sizes of the buffer after each read of the cursor (``_fill``), as the result set's own count of replies.""" + seen = [] + real = database.NativeResultSet._fill + + def counting(self): + got = real(self) + seen.append(len(self._buf)) + return got + + monkeypatch.setattr(database.NativeResultSet, "_fill", counting) + return seen + + +@pytest.mark.parametrize("size", [1, 7, 255, 256, 257, 2048, 2049, 4999, 5000, 5001]) +def test_batches_equal_to_json_list_cut_at_batch_size(db, size): + want = _select(db).to_json_list() + assert len(want) == ROWS + got = list(_select(db).iter_json_batches(batch_size=size)) + assert [len(b) for b in got] == [size] * (ROWS // size) + ( + [ROWS % size] if ROWS % size else [] + ) + assert [row for batch in got for row in batch] == want + assert want[0] == {"id": 0, "txt": "t0", "val": 0.0} + assert want[-1] == {"id": ROWS - 1, "txt": f"t{ROWS - 1}", "val": (ROWS - 1) / 8.0} + + +def test_default_batch_size_and_an_empty_result(db): + rows = [r for b in _select(db).iter_json_batches() for r in b] + assert [r["id"] for r in rows] == list(range(ROWS)) + assert list(db.query("sql", "SELECT FROM R WHERE id < 0").iter_json_batches()) == [] + + +def test_the_last_batch_finishes_the_cursor(db): + rs = db.query("sql", "SELECT id FROM R WHERE id < 10 ORDER BY id") + (batch,) = list(rs.iter_json_batches(batch_size=10)) # exactly full, and the last + assert len(batch) == 10 + assert ( + list(rs.iter_json_batches()) == [] and rs.to_list() == [] + ) # read to the end, not closed early + + +def test_the_buffer_holds_one_reply_not_the_result(db, fills): + rs = _select(db) + held = [] + for batch in rs.iter_json_batches(batch_size=100): + # rows in hand: this batch plus what the result set still buffers + held.append(len(batch) + len(rs._buf)) + assert fills, "the cursor was never read past its first reply" + assert max(fills) <= database._NEXT_BATCH + assert max(held) <= 100 + database._NEXT_BATCH + assert max(held) < ROWS // 2 + + +def test_the_cursor_is_read_when_the_next_batch_is_asked_for(db, fills): + rs = _select(db) + batches = rs.iter_json_batches(batch_size=1000) + first = next(batches) + assert len(first) == 1000 + # the first batch needs one read past the first reply of 256 rows, not all of the 5000 rows + assert len(fills) <= 1 + assert rs._more and not rs._closed + rest = [r for b in batches for r in b] + assert [r["id"] for r in first + rest] == list(range(ROWS)) + assert rs._exhausted and rs._closed + + +def test_close_between_two_batches_raises_and_closes_the_cursor(db): + rs = _select(db) + batches = rs.iter_json_batches(batch_size=100) + assert len(next(batches)) == 100 + rs.close() + assert rs._closed and not rs._more and not rs._buf + with pytest.raises(arcadedb.ArcadeDBError, match="closed before"): + next(batches) + + +def test_leaving_the_with_block_mid_iteration_closes_the_cursor(db): + with _select(db) as rs: + batches = rs.iter_json_batches(batch_size=100) + assert len(next(batches)) == 100 + assert rs._closed and not rs._exhausted + with pytest.raises(arcadedb.ArcadeDBError, match="closed before"): + rs.to_list() + + +def test_an_abandoned_generator_leaves_the_result_to_its_owner(db): + rs = _select(db) + batches = rs.iter_json_batches(batch_size=100) + next(batches) + batches.close() # the generator is gone, the cursor is the result set's + assert not rs._closed + rs.close() + assert rs._closed + + +def test_a_batch_size_below_one_is_refused_when_the_first_batch_is_asked_for(db): + batches = _select(db).iter_json_batches(batch_size=0) + with pytest.raises(ValueError, match="at least 1"): + next(batches) From 255a7b343c4b85ca0864363401c13eac8d8d539b Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Wed, 4 Nov 2026 13:53:48 +0900 Subject: [PATCH 2/2] iter_json_batches: docs say batch_size bounds memory on native; buffer-bound test counts rows read ahead so it fails on main The old test watched the result set's buffer, which the drain on main empties as it goes, so it passed there. The new one counts rows read from the cursor against rows delivered at every batch: main has read all 5000 after the first batch of 100, the change at most one reply ahead. Docs carry the measured ratios with 95% intervals from the laptop study (copied by script from ab.txt). Co-Authored-By: Claude Opus 5.5 --- .../pages/how-to/databases-and-queries.adoc | 3 ++ src/arcadedb/_native/database.py | 7 +++++ src/arcadedb/results.py | 6 ++++ tests/test_native_streaming_batches.py | 31 +++++++++++++------ 4 files changed, 38 insertions(+), 9 deletions(-) diff --git a/docs/modules/ROOT/pages/how-to/databases-and-queries.adoc b/docs/modules/ROOT/pages/how-to/databases-and-queries.adoc index 3b372fd..04eddc9 100644 --- a/docs/modules/ROOT/pages/how-to/databases-and-queries.adoc +++ b/docs/modules/ROOT/pages/how-to/databases-and-queries.adoc @@ -238,6 +238,9 @@ A query vector for `vectorNeighbors` can be a numpy array or a list. The convers * `first()` returns the first row, or `None` when there is none, so `if db.query(...).first():` tests whether anything matched. * `to_list()` materialises the whole result as a list of dicts. Use it when you need the whole result in memory; iterate when you stop early or the result is large. +* `iter_json_batches(batch_size=10_000)` yields the rows as lists of JSON-native dicts, `batch_size` rows per list (the last may be shorter), with the same values as `to_json_list()`. On the native build `batch_size` bounds memory: the result set holds one reply of the cursor (2048 rows) and one batch at a time, and reads the next reply only when the next batch is asked for, so a large result is never held whole. On the JVM build a batch is serialised on the Java side one at a time. `to_list()`, `to_json_list()`, `to_dataframe()`, `to_columns()`, `to_arrow()` and `to_polars()` still hold the whole result, because they return it whole. ++ +One laptop study on the native build, 1,000,000 rows read with `batch_size=10_000`, 10 clean pairs of this change against the previous behaviour; each figure is the ratio of the new value to the old one, so below 1 is lower. Peak RSS: 0.286 (95% interval 0.284 to 0.289). Time to the first batch: 0.010 (95% interval 0.009 to 0.010). Wall time of the whole read: 0.958 (95% interval 0.952 to 0.966). This is one machine and one result shape, not a speed promise. === Record ids in rows diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index 865dd11..d6a3b2b 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -326,6 +326,13 @@ def to_json_list(self, batch_size: int = 10_000) -> List[Dict[str, Any]]: return [json_native(row) for row in self._drain_dicts()] def iter_json_batches(self, batch_size: int = 10_000, **kwargs): + """Yield the rows as lists of JSON-native dicts, ``batch_size`` rows per list (the last may be shorter). + + The values are those of :meth:`to_json_list`. ``batch_size`` bounds memory: one reply of the cursor (``_NEXT_BATCH`` rows) + and one batch are held, and the cursor is read when the next batch is asked for. + + One laptop study on the native build, 1,000,000 rows read with ``batch_size=10_000``, 10 clean pairs of this change against the previous behaviour; each figure is the ratio of the new value to the old one, so below 1 is lower. Peak RSS: 0.286 (95% interval 0.284 to 0.289). Time to the first batch: 0.010 (95% interval 0.009 to 0.010). Wall time of the whole read: 0.958 (95% interval 0.952 to 0.966). This is one machine and one result shape, not a speed promise. + """ if int(batch_size) < 1: raise ValueError("batch_size must be at least 1") for batch in self._iter_dict_batches(int(batch_size)): diff --git a/src/arcadedb/results.py b/src/arcadedb/results.py index 4c00cd9..7c491e3 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -609,6 +609,12 @@ def iter_json_batches( Streaming counterpart of :meth:`to_json_list` with the same JSON-native type semantics; bounds memory to one batch. Falls back to chunked per-row conversion when the bridge jar is unavailable. + + On the native build the same property holds: ``batch_size`` bounds + memory (one reply of the cursor plus one batch), and the next reply is + read only when the next batch is asked for. + + One laptop study on the native build, 1,000,000 rows read with ``batch_size=10_000``, 10 clean pairs of this change against the previous behaviour; each figure is the ratio of the new value to the old one, so below 1 is lower. Peak RSS: 0.286 (95% interval 0.284 to 0.289). Time to the first batch: 0.010 (95% interval 0.009 to 0.010). Wall time of the whole read: 0.958 (95% interval 0.952 to 0.966). This is one machine and one result shape, not a speed promise. """ if int(batch_size) < 1: raise ValueError(f"batch_size must be at least 1, got {batch_size}") diff --git a/tests/test_native_streaming_batches.py b/tests/test_native_streaming_batches.py index e864bfb..5176b4b 100644 --- a/tests/test_native_streaming_batches.py +++ b/tests/test_native_streaming_batches.py @@ -33,13 +33,14 @@ def _select(db): @pytest.fixture def fills(monkeypatch): - """The sizes of the buffer after each read of the cursor (``_fill``), as the result set's own count of replies.""" + """The number of rows each read of the cursor (``_fill``) added to the result set's buffer, one entry per read.""" seen = [] real = database.NativeResultSet._fill def counting(self): + before = len(self._buf) got = real(self) - seen.append(len(self._buf)) + seen.append(len(self._buf) - before) return got monkeypatch.setattr(database.NativeResultSet, "_fill", counting) @@ -74,16 +75,28 @@ def test_the_last_batch_finishes_the_cursor(db): ) # read to the end, not closed early -def test_the_buffer_holds_one_reply_not_the_result(db, fills): +def test_rows_read_ahead_stay_within_one_reply_not_the_result(db, fills): + """Rows read from the cursor but not yet handed to the caller stay within one reply plus one batch. + + The old code drained the whole result into dicts before it cut the first batch, so after the first batch it had read all + ROWS rows and the caller held 100: 4900 rows ahead. The buffer alone cannot show that (the drain empties it as it goes), so + this counts the rows read from the cursor against the rows delivered, at every batch. + """ rs = _select(db) - held = [] - for batch in rs.iter_json_batches(batch_size=100): - # rows in hand: this batch plus what the result set still buffers - held.append(len(batch) + len(rs._buf)) + size = 100 + read = len(rs._buf) # the first reply, read when the query ran + delivered = 0 + ahead = [] + for batch in rs.iter_json_batches(batch_size=size): + delivered += len(batch) + ahead.append(read + sum(fills) - delivered) + assert delivered == ROWS and read + sum(fills) == ROWS + assert len(ahead) == ROWS // size assert fills, "the cursor was never read past its first reply" assert max(fills) <= database._NEXT_BATCH - assert max(held) <= 100 + database._NEXT_BATCH - assert max(held) < ROWS // 2 + assert ahead[0] <= database._NEXT_BATCH # not ROWS - size + assert max(ahead) <= database._NEXT_BATCH + size + assert max(ahead) < ROWS // 2 def test_the_cursor_is_read_when_the_next_batch_is_asked_for(db, fills):