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
63 changes: 63 additions & 0 deletions Cargo.lock

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

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ pyo3 = { version = "0.29" }
# here is the C Data Interface, which is how a result reaches pyarrow,
# pandas and polars without a Python object per cell.
arrow = { version = "59", default-features = false, features = ["ffi"] }
numpy = "0.29"

# Which ABI to build against, which is the one thing a wheel cannot be
# told after it is built. The default is the stable ABI from 3.11 up,
Expand Down
5 changes: 4 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -136,10 +136,13 @@ 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.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.

`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.

## Reading a result as it arrives

`execute` runs the statement to the end and hands back every row. `stream` hands back the rows as the engine makes them, which is what you want when the result is bigger than the memory you meant to spend on it, or when the first rows are worth having before the last are made.
Expand Down Expand Up @@ -353,7 +356,7 @@ Half of this package is compiled, which is the one thing an inspection cannot se

## What works today

The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch a row or a block at a time, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, an appender for growing one, transactions as a context manager that commits at the end of a block and rolls back when it raises, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, `register` for putting a frame under a name a statement can match on and reading it where it lies, stubs inside the wheel with a gate that keeps them true, the GIL released around every statement, every load and every copy out, `Ctrl-C` and `interrupt()` stopping a statement without touching the connection under it, `zudb.aio` for the same calls awaited on an event loop, results, nodes, rels and paths that draw themselves in a notebook with `%gql` and `%%gql` to run statements in one, `zudb.dbapi` for code written against PEP 249, `prepare`, `explain` and `profile` for a statement compiled once, the plan it would run and the plan it did, `stream` for a result read as the engine makes it rather than after it has made all of it, and `cursor()` for a second connection made from the first, which is how a pool is written. Each one landed with the tests that say it works.
The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch a row or a block at a time, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, an appender for growing one, transactions as a context manager that commits at the end of a block and rolls back when it raises, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns, as numpy arrays, and as pandas and polars frames, `register` for putting a frame under a name a statement can match on and reading it where it lies, stubs inside the wheel with a gate that keeps them true, the GIL released around every statement, every load and every copy out, `Ctrl-C` and `interrupt()` stopping a statement without touching the connection under it, `zudb.aio` for the same calls awaited on an event loop, results, nodes, rels and paths that draw themselves in a notebook with `%gql` and `%%gql` to run statements in one, `zudb.dbapi` for code written against PEP 249, `prepare`, `explain` and `profile` for a statement compiled once, the plan it would run and the plan it did, `stream` for a result read as the engine makes it rather than after it has made all of it, and `cursor()` for a second connection made from the first, which is how a pool is written. Each one landed with the tests that say it works.

## Wheels

Expand Down
7 changes: 6 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,12 @@ Issues = "https://github.com/tamnd/zu-python/issues"
arrow = ["pyarrow>=14"]
pandas = ["pandas>=2.0", "pyarrow>=14"]
polars = ["polars>=1.3"]
all = ["pyarrow>=14", "pandas>=2.0", "polars>=1.3"]
# numpy 1.23 is where `numpy.dtypes` and the 1.x ABI settled, and it is
# older than anything a caller is likely to be holding. Nothing else
# comes with it: `fetchnumpy` wraps the engine's own buffers and needs
# no Arrow on the way.
numpy = ["numpy>=1.23"]
all = ["pyarrow>=14", "pandas>=2.0", "polars>=1.3", "numpy>=1.23"]

[dependency-groups]
dev = [
Expand Down
3 changes: 3 additions & 0 deletions python/zudb/_zudb.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -477,6 +477,9 @@ class Result:
def to_pandas(self) -> Any:
"""The rows as a `pandas.DataFrame`, with Arrow-backed dtypes."""

def fetchnumpy(self) -> dict[str, Any]:
"""The columns as numpy arrays, in a dict keyed by column name."""

def to_polars(self) -> Any:
"""The rows as a `polars.DataFrame`."""

Expand Down
26 changes: 26 additions & 0 deletions src/conn.rs
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use crate::columns;
use crate::error::{closed, programming, to_py_err};
use crate::html;
use crate::interrupt;
use crate::numpy;
use crate::plan;
use crate::prepared::Prepared;
use crate::register;
Expand Down Expand Up @@ -823,6 +824,31 @@ impl Result {
Self::to_arrow(slf)?.call_method("to_pandas", (), Some(&how))
}

/// The columns as numpy arrays, in a dict keyed by column name.
///
/// For the code that takes arrays rather than frames, and it needs
/// numpy and nothing else: the engine's own buffers are what numpy
/// wraps, so an integer, float, datetime or duration column is a
/// move rather than a copy and pyarrow never comes into it.
///
/// A column with a null in it comes back as a
/// `numpy.ma.masked_array`, since numpy has no missing integer.
/// Strings, nodes, rels, paths, lists and records come back as
/// object arrays, which have somewhere to put a `None` and hold a
/// Python object per row like the rows themselves do. Dates are
/// `datetime64[D]`, datetimes `datetime64[ns]`, 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 is the closest thing numpy has to one.
///
/// Reading it does not move the cursor `fetchone` uses, like every
/// other way of reading the columns: a caller who took a copy of
/// the result has not taken any of its rows.
fn fetchnumpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyDict>> {
needed(py, "numpy", "numpy")?;
numpy::arrays(py, &self.result, &self.names)
}

/// The rows as a `polars.DataFrame`.
///
/// The constructor rather than `from_arrow`, because `from_arrow`
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ mod frame;
mod html;
mod interrupt;
mod load;
mod numpy;
mod plan;
mod prepared;
mod register;
Expand Down
Loading
Loading