From 74c10e9822f23b325b33f4c7809df66a04333713 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Wed, 4 Nov 2026 06:48:32 +0900 Subject: [PATCH 1/3] 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 f8f65645..668e3425 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -291,6 +291,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() @@ -314,9 +344,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 9722c2d1..03f41461 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -67,6 +67,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 00000000..e864bfb0 --- /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 9d8f21718ba4ce0e0ae41dc694865f214d795f02 Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Wed, 4 Nov 2026 13:53:48 +0900 Subject: [PATCH 2/3] 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 ba9d927d..e685d95e 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 668e3425..e0af5204 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -342,6 +342,13 @@ def to_json_list(self, batch_size: int = 10_000) -> List[Dict[str, Any]]: return [json_native(row) for row in rows] 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 b30dfab7..c2e11377 100644 --- a/src/arcadedb/results.py +++ b/src/arcadedb/results.py @@ -617,6 +617,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 e864bfb0..5176b4ba 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): From 6cf8f55e87abf865bb4cecd683eb45536b612c9a Mon Sep 17 00:00:00 2001 From: Taewoon Kim Date: Thu, 5 Nov 2026 06:44:46 +0900 Subject: [PATCH 3/3] iter_json_batches: skip the json_native walk for a result with no tagged reply (as to_json_list after #223) A batch is built after the cursor reads the replies its rows came from, so the sticky _tagged flag read at that moment covers every row of the batch: a result that never restored a tagged value is yielded as parsed. Tests: no walk for an untagged result, and a DATE that first appears in a later reply is still converted at two batch sizes (equal to to_json_list). Co-Authored-By: Claude Opus 5.5 --- src/arcadedb/_native/database.py | 4 ++- tests/test_native_streaming_batches.py | 49 ++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) diff --git a/src/arcadedb/_native/database.py b/src/arcadedb/_native/database.py index e0af5204..edf5cc30 100644 --- a/src/arcadedb/_native/database.py +++ b/src/arcadedb/_native/database.py @@ -352,7 +352,9 @@ 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") for batch in self._iter_dict_batches(int(batch_size)): - yield [json_native(row) for row in batch] + # as to_json_list: no reply read so far restored a tagged value, so every value is a json.loads type that json_native + # returns as it is. The flag is read after the batch is built, so it covers every reply the batch's rows came from + yield [json_native(row) for row in batch] if self._tagged else batch def to_dataframe(self, convert_types: bool = True): import pandas as pd diff --git a/tests/test_native_streaming_batches.py b/tests/test_native_streaming_batches.py index 5176b4ba..b3628f19 100644 --- a/tests/test_native_streaming_batches.py +++ b/tests/test_native_streaming_batches.py @@ -7,6 +7,8 @@ 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. """ +from datetime import date + import pytest import arcadedb @@ -145,3 +147,50 @@ 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) + + +@pytest.fixture +def walked(monkeypatch): + """The number of rows ``json_native`` (the walk that converts DATE, DATETIME, DECIMAL, sets and record ids) was called for.""" + seen = [0] + real = database.json_native + + def counting(value): + seen[0] += 1 + return real(value) + + monkeypatch.setattr(database, "json_native", counting) + return seen + + +def test_batches_of_an_untagged_result_are_not_walked(db, walked): + got = [r for b in _select(db).iter_json_batches(batch_size=100) for r in b] + assert len(got) == ROWS + assert ( + walked[0] == 0 + ), "no reply was tagged, so iter_json_batches must skip json_native like to_json_list does" + assert got == _select(db).to_json_list() + + +@pytest.mark.parametrize("size", [100, 1000]) +def test_a_tag_in_a_later_reply_is_converted_in_every_batch_it_reaches( + temp_db, walked, size +): + """The first reply (256 rows) has no tagged value and a later one does: its rows must still arrive as JSON-native values.""" + temp_db.command("sql", "CREATE DOCUMENT TYPE D") + temp_db.command("sql", "CREATE PROPERTY D.id INTEGER") + temp_db.command("sql", "CREATE PROPERTY D.d DATE") + rows = [{"id": i} for i in range(600)] + for row in rows[300:]: + row["d"] = date(2026, 1, 2) + temp_db.insert_many("D", rows) + query = "SELECT id, d FROM D ORDER BY id" + want = temp_db.query("sql", query).to_json_list() + walked[0] = 0 + got = [r for b in temp_db.query("sql", query).iter_json_batches(size) for r in b] + assert got == want + assert want[0].get("d") is None + assert ( + type(want[599]["d"]) is int + ), "the DATE is epoch milliseconds, not a date object" + assert walked[0] > 0