Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 21 additions & 10 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

17 changes: 11 additions & 6 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -18,17 +18,22 @@ crate-type = ["cdylib"]
# with (ADR 0002), so a revision is the honest way to say which one.
# A local checkout is used instead with a `paths` override in
# `.cargo/config.toml`, which is untracked on purpose.
zudb = { package = "zu", git = "https://github.com/tamnd/zu", rev = "130f67db924bcd0f766ee814b0da2edae32150d4" }
zu-common = { git = "https://github.com/tamnd/zu", rev = "130f67db924bcd0f766ee814b0da2edae32150d4" }
zudb = { package = "zu", git = "https://github.com/tamnd/zu", rev = "0698a4eccd31670f0a875b6d097d7753440a51b3" }
zu-common = { git = "https://github.com/tamnd/zu", rev = "0698a4eccd31670f0a875b6d097d7753440a51b3" }
# The one translation from a result into Arrow, which lives in the
# engine tree so that every client agrees about what a column becomes.
# `ffi` is the only feature this client turns on: what Python wants is
# the C Data Interface, which is how a result reaches pyarrow, pandas
# and polars without a Python object per cell.
zu-arrow = { git = "https://github.com/tamnd/zu", rev = "0698a4eccd31670f0a875b6d097d7753440a51b3", features = ["ffi"] }
# `extension-module` is asked for by maturin, in pyproject.toml, and
# not here. Only the build backend knows how an extension is linked on
# the platform it is building for, and a crate that turns the feature
# on by default is a crate `cargo build` cannot link on its own.
pyo3 = { version = "0.29" }
# Arrow, for the columns a result leaves as. Only `ffi` is asked for:
# the readers and writers are the engine's business and what is wanted
# here is the C Data Interface, which is how a result reaches pyarrow,
# pandas and polars without a Python object per cell.
# Arrow again, directly, because reading one is this client's job too:
# `register` takes a frame from pandas or polars over the same C Data
# Interface, and that side is a reader and not a writer.
arrow = { version = "59", default-features = false, features = ["ffi"] }
numpy = "0.29"

Expand Down
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -135,13 +135,13 @@ result = conn.execute("MATCH (p:person) RETURN p.name AS name, p.score AS score"
result.to_arrow() # pyarrow.Table
result.to_pandas() # DataFrame with Arrow-backed dtypes
result.to_polars() # polars.DataFrame
result.record_batches() # a reader, for a result larger than memory
result.record_batches() # a reader, for a consumer that writes as it reads
result.fetchnumpy() # {"name": array([...]), "score": array([...])}
```

`Result` implements `__arrow_c_stream__`, so anything that reads the protocol reads a result directly and none of the four methods above is needed: `pyarrow.table(result)` and `polars.DataFrame(result)` both work. Batches are 65,536 rows. A column holds one type, which the values decide, and integers beside floats are the one mixture that widens rather than being refused. Nodes, rels and paths go across as structs. The copy runs with the GIL released, and on this machine 300,000 rows across three columns take 44 ms as Arrow against 67 ms as Python objects, and a single integer column takes 13.8 ms against 44.5 ms.
`Result` implements `__arrow_c_stream__`, so anything that reads the protocol reads a result directly and none of the four methods above is needed: `pyarrow.table(result)` and `polars.DataFrame(result)` both work. Batches are 65,536 rows, or whatever `record_batches(rows_per_batch)` asks for, and the size costs nothing either way: the arrays are built once and a batch is a slice of them rather than a copy. A column holds one type, which the values decide, and integers beside floats are the one mixture that widens rather than being refused. A result that matched no rows still says what its columns hold, so a query that found nothing can still be written to Parquet or appended to a table that already exists. Nodes, rels and paths go across as structs. The translation runs with the GIL released, and on this machine 300,000 rows across three columns take 4.8 ms as Arrow against 86 ms as Python objects, and a single integer column takes 0.8 ms against 43 ms. That is a wider gap than it used to be because the engine now fills the column buffers during the scan rather than transposing the rows afterwards, so what is left here is putting an Arrow type around a buffer that already exists.

`fetchnumpy()` is the same columns as numpy arrays, keyed by column name, for the code that takes arrays rather than frames. It needs numpy and nothing else: an integer, float, datetime or duration column is the engine's own buffer moved into numpy and named, so there is no pass over the values and no second copy of the result in memory. Over a million rows in two columns on this machine it takes 26 ms, against 47 ms for building the Arrow table and calling `to_numpy` on each of its columns, and 130 ms for the same rows as tuples. Dates are `datetime64[D]`, datetimes `datetime64[ns]`, and durations and times of day `timedelta64[ns]`, the last of those being nanoseconds since midnight, which is what a clock reading is on a number line and the closest thing numpy has to one. A column with a null in it comes back as a `numpy.ma.masked_array`, since numpy has no missing integer, and the mask is built from the validity bitmap the engine already filled rather than by walking the column again. Strings, nodes, rels, paths, lists and records come back as object arrays with `None` in the cell, because an object array has somewhere to put one.
`fetchnumpy()` is the same columns as numpy arrays, keyed by column name, for the code that takes arrays rather than frames. It needs numpy and nothing else: an integer, float, datetime or duration column is the engine's own buffer moved into numpy and named, so there is no pass over the values and no second copy of the result in memory. Over a million rows in two columns on this machine it takes 9.9 ms, against 11.9 ms for building the Arrow table and calling `to_numpy` on each of its columns, and 199 ms for the same rows as tuples. Dates are `datetime64[D]`, datetimes `datetime64[ns]`, and durations and times of day `timedelta64[ns]`, the last of those being nanoseconds since midnight, which is what a clock reading is on a number line and the closest thing numpy has to one. A column with a null in it comes back as a `numpy.ma.masked_array`, since numpy has no missing integer, and the mask is built from the validity bitmap the engine already filled rather than by walking the column again. Strings, nodes, rels, paths, lists and records come back as object arrays with `None` in the cell, because an object array has somewhere to put one.

## Reading a result as it arrives

Expand Down
6 changes: 5 additions & 1 deletion conformance/cases.py
Original file line number Diff line number Diff line change
Expand Up @@ -217,7 +217,11 @@ def _case(node: Node) -> Case:
f"line {line}: a case says what it produces, with `columns:` and `rows:` or with "
"`raises:`"
)
names = columns_node.seq()
# Empty counts, because `FINISH` is a query that answers no columns
# at all, which is not the same as a query whose columns held no
# rows, and the corpus writes it as a `columns:` with nothing under
# it.
names = columns_node.seq_or_empty()
if names is None:
raise CorpusError(f"line {line}: `columns:` is a sequence of names")
columns = []
Expand Down
16 changes: 15 additions & 1 deletion conformance/reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -377,7 +377,21 @@ def _scalar(text: str, line: int) -> Node:
return Node("scalar", line, text=text, quoted=False)


_ESCAPES = {'"': '"', "\\": "\\", "n": "\n", "r": "\r", "t": "\t", "0": "\0"}
# The escapes the corpus uses, which is a subset of YAML's. The ones
# that name a code point by its digits are not here, because the corpus
# writes those as the character itself and a case that wants the digits
# is testing the engine's own escapes inside a query rather than the
# file's.
_ESCAPES = {
'"': '"',
"\\": "\\",
"n": "\n",
"r": "\r",
"t": "\t",
"0": "\0",
"b": "\b",
"f": "\f",
}


def _unescape(body: str, line: int) -> str:
Expand Down
2 changes: 1 addition & 1 deletion python/zudb/_zudb.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -483,7 +483,7 @@ class Result:
def to_polars(self) -> Any:
"""The rows as a `polars.DataFrame`."""

def record_batches(self) -> Any:
def record_batches(self, rows_per_batch: int | None = None) -> Any:
"""The rows as a `pyarrow.RecordBatchReader`, a batch at a time."""

def __arrow_c_stream__(self, requested_schema: object | None = None) -> Any:
Expand Down
Loading
Loading