diff --git a/Cargo.lock b/Cargo.lock index 8a04620..519ca82 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -822,12 +822,37 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "matrixmultiply" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7" +dependencies = [ + "autocfg", + "rawpointer", +] + [[package]] name = "memchr" version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -875,6 +900,22 @@ dependencies = [ "libm", ] +[[package]] +name = "numpy" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + [[package]] name = "object_store" version = "0.14.1" @@ -951,6 +992,15 @@ version = "1.15.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85" +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.6" @@ -1041,6 +1091,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1093,6 +1149,12 @@ dependencies = [ "smallvec", ] +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + [[package]] name = "rustc_version" version = "0.4.1" @@ -1752,6 +1814,7 @@ name = "zudb-python" version = "0.0.1" dependencies = [ "arrow", + "numpy", "pyo3", "zu", "zu-common", diff --git a/Cargo.toml b/Cargo.toml index 7a8715a..1c9021c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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, diff --git a/README.md b/README.md index b611812..01665d2 100644 --- a/README.md +++ b/README.md @@ -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. @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 7de05f5..6d70b1e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = [ diff --git a/python/zudb/_zudb.pyi b/python/zudb/_zudb.pyi index bd4c557..35c3d1d 100644 --- a/python/zudb/_zudb.pyi +++ b/python/zudb/_zudb.pyi @@ -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`.""" diff --git a/src/conn.rs b/src/conn.rs index a22fb93..90f2a92 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -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; @@ -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> { + 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` diff --git a/src/lib.rs b/src/lib.rs index 71997ef..4c89a7b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,7 @@ mod frame; mod html; mod interrupt; mod load; +mod numpy; mod plan; mod prepared; mod register; diff --git a/src/numpy.rs b/src/numpy.rs new file mode 100644 index 0000000..1ec865a --- /dev/null +++ b/src/numpy.rs @@ -0,0 +1,237 @@ +//! A result as numpy arrays, one per column. +//! +//! The third way out of a result, beside rows and Arrow. It exists +//! because numpy is what a lot of code already holds: a model that +//! takes arrays, a plotting call, a loop somebody wrote before +//! DataFrames. Handing that code an Arrow table means it has to convert +//! one, and converting one means pyarrow has to be installed to do +//! work numpy could have done with the same bytes. +//! +//! The same bytes is what this is. `zudb::query::column` hands back one +//! owned buffer per column, values end to end in the layout every +//! columnar format uses, and numpy is a pointer, a length and a dtype +//! over exactly that. So an integer column, a float column, a datetime +//! column and a duration column all become arrays by moving the `Vec` +//! into numpy and naming the type: no pass over the values, no +//! allocation, and no second copy of the result in memory. The two that +//! cost something are the ones where the layouts differ rather than the +//! names: a boolean column is a bit per row here and a byte per row +//! there, and a date is 32 bits here and 64 bits there, so each takes +//! one widening pass. +//! +//! Nulls are the other half of the shape. numpy has no missing value +//! for an integer, so a column with one comes back as a +//! `numpy.ma.masked_array`, which is a data array and a boolean mask +//! beside it and is what a masked column has meant in numpy since long +//! before any of the alternatives. The mask is built from the validity +//! bitmap the engine already filled, and the data array underneath is +//! still the engine's own buffer: masking is a wrapper and not a copy. +//! Columns that become object arrays carry `None` in the cell instead, +//! because an object array has somewhere to put it. + +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use zudb::query::QueryResult; +use zudb::query::column::{Column, ColumnData, ColumnType, Offsets, Validity}; + +// The crate, not this module. Both are called numpy and a plain path +// would be ambiguous, which is the price of naming a module after the +// thing it is about. +use ::numpy::IntoPyArray; + +use crate::value::{Names, to_py}; + +/// Every column of a result, as a dict of numpy arrays. +/// +/// The order is the order the statement projected them, which a dict +/// keeps, so `list(result.fetchnumpy())` is `result.columns`. +pub fn arrays<'py>( + py: Python<'py>, + result: &QueryResult, + names: &Names, +) -> PyResult> { + let ma = PyModule::import(py, "numpy.ma")?; + // The read down the columns touches no Python object, and it is + // two passes over every row of the result, so it happens with the + // GIL down like every other pass this client makes over a result. + let columns = py + .detach(|| result.columnar()) + .map_err(|mixed| PyTypeError::new_err(mixed.to_string()))?; + + let out = PyDict::new(py); + for held in columns.columns { + let name = held.name; + if out.contains(name)? { + return Err(PyValueError::new_err(format!( + "the result has two columns called '{name}', and a dict holds one of each name: \ + give them different names with AS" + ))); + } + let array = column(py, &ma, held, names)?; + out.set_item(name, array)?; + } + Ok(out) +} + +/// One column, as the array and the mask that go with it. +fn column<'py>( + py: Python<'py>, + ma: &Bound<'py, PyModule>, + held: Column<'_>, + names: &Names, +) -> PyResult> { + let name = held.name; + let len = held.len; + let valid = held.validity; + + // The object arms answer for themselves: an object array has a cell + // for `None`, so a mask beside it would say the same thing twice. + let flat = match held.data { + ColumnData::Null => return objects(py, (0..len).map(|_| py.None()).collect()), + ColumnData::Str(strings) => { + let mut out = Vec::with_capacity(len); + let bytes = &strings.bytes; + let mut spans = spans(&strings.offsets); + for at in 0..len { + let (from, upto) = spans.next().unwrap_or((0, 0)); + out.push(match missing(&valid, at) { + true => py.None(), + false => text(py, name, at, &bytes[from..upto])?, + }); + } + return objects(py, out); + } + ColumnData::Complex(values) => { + let mut out = Vec::with_capacity(values.len()); + for value in values { + out.push(to_py(py, value, names)?.unbind()); + } + return objects(py, out); + } + + // A bit per row here, a byte per row in numpy, so this one is + // unpacked rather than moved. + ColumnData::Bool { bits } => { + let mut out = Vec::with_capacity(len); + for at in 0..len { + out.push(bits[at / 8] & (1u8 << (at % 8)) != 0); + } + out.into_pyarray(py).into_any() + } + ColumnData::Int(values) => values.into_pyarray(py).into_any(), + ColumnData::Float(values) => values.into_pyarray(py).into_any(), + // Days are 32 bits in the engine and 64 in a numpy + // `datetime64`, which is the one widening pass here. + ColumnData::Days(values) => { + let wide: Vec = values.into_iter().map(i64::from).collect(); + seen(wide.into_pyarray(py).into_any(), "datetime64[D]")? + } + ColumnData::Months(values) => seen(values.into_pyarray(py).into_any(), "timedelta64[M]")?, + ColumnData::Nanos(values) => { + let array = values.into_pyarray(py).into_any(); + match held.ty { + // A time of day is nanoseconds since midnight, which is + // what numpy calls a `timedelta64`. It has no type for a + // clock reading, and this is the reading itself rather + // than a rounded stand-in for it. + ColumnType::LocalTime => seen(array, "timedelta64[ns]")?, + ColumnType::DayTime => seen(array, "timedelta64[ns]")?, + ColumnType::LocalDatetime => seen(array, "datetime64[ns]")?, + // The instant, in UTC. numpy has no zone to carry the + // offset in, and the instant is what the buffer holds. + ColumnType::ZonedDatetime { .. } => seen(array, "datetime64[ns]")?, + ColumnType::ZonedTime { .. } => { + return Err(PyTypeError::new_err(format!( + "column '{name}' holds a time with an offset, which numpy has no type for" + ))); + } + ty => return Err(mismatch(name, &ty)), + } + } + }; + + match valid { + None => Ok(flat), + Some(held) => masked(ma, flat, mask(&held, len).into_pyarray(py).into_any()), + } +} + +/// The bytes of one string, as a Python one. +/// +/// The engine wrote the buffer and every string in it came in as a +/// Python string, so this cannot fail in practice; it is checked +/// anyway, because the alternative is a wrong answer rather than an +/// exception. +fn text<'py>(py: Python<'py>, name: &str, at: usize, bytes: &[u8]) -> PyResult> { + match std::str::from_utf8(bytes) { + Ok(text) => Ok(pyo3::types::PyString::new(py, text).into_any().unbind()), + Err(_) => Err(PyValueError::new_err(format!( + "the string at row {at} of column '{name}' is not valid UTF-8" + ))), + } +} + +/// Where each string sits in the bytes, whichever width the offsets are. +fn spans(offsets: &Offsets) -> Box + '_> { + match offsets { + Offsets::I32(held) => Box::new( + held.windows(2) + .map(|pair| (pair[0] as usize, pair[1] as usize)), + ), + Offsets::I64(held) => Box::new( + held.windows(2) + .map(|pair| (pair[0] as usize, pair[1] as usize)), + ), + } +} + +/// Whether row `at` is null, when there is a bitmap to ask. +fn missing(valid: &Option, at: usize) -> bool { + valid.as_ref().is_some_and(|held| !held.is_valid(at)) +} + +/// The mask numpy wants, which is the bitmap the other way up: set +/// means missing there, and set means present in the engine's. +fn mask(valid: &Validity, len: usize) -> Vec { + (0..len).map(|at| !valid.is_valid(at)).collect() +} + +/// A list of Python objects as an object array. +fn objects<'py>(py: Python<'py>, values: Vec>) -> PyResult> { + Ok(values.into_pyarray(py).into_any()) +} + +/// The same buffer read as another type of the same width. +/// +/// `view` is numpy's own word for it and copies nothing: a +/// `datetime64[ns]` array and the `int64` array under it are the same +/// bytes, which is exactly the relationship the engine's buffer already +/// has with both of them. +fn seen<'py>(array: Bound<'py, PyAny>, dtype: &str) -> PyResult> { + array.call_method1("view", (dtype,)) +} + +/// The array and its mask, as the one object numpy has for the pair. +fn masked<'py>( + ma: &Bound<'py, PyModule>, + data: Bound<'py, PyAny>, + mask: Bound<'py, PyAny>, +) -> PyResult> { + let how = PyDict::new(ma.py()); + // Neither array is copied: the data is the engine's buffer and the + // mask was built for this call, so there is nothing to protect + // either of them from. + how.set_item("copy", false)?; + ma.call_method("masked_array", (data, mask), Some(&how)) +} + +/// A buffer holding something other than what the column's type says, +/// which is this module reading its own input wrong rather than +/// anything the caller did. +fn mismatch(name: &str, ty: &ColumnType) -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err(format!( + "column '{name}' came back as {} in a buffer that does not hold one", + ty.name() + )) +} diff --git a/tests/test_import.py b/tests/test_import.py index cfe1392..d754668 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -132,6 +132,27 @@ def test_the_dbapi_module_is_not_imported_by_importing_this_one() -> None: assert done.stdout.strip() == "False" +def test_numpy_arrives_when_a_result_is_asked_for_its_arrays(tmp_path: Path) -> None: + """The extension links numpy's C API and still does not import it. + + It is loaded at the first array and not at module init, which is the + difference between a client that costs numpy to import and one that + costs it to use. + """ + pytest.importorskip("numpy") + path = tmp_path / "arrays.zu1" + done = run( + "import sys, zudb\n" + f"zudb.load({str(path)!r}, nodes='person', columns={{'uid': [1, 2, 3]}})\n" + f"conn = zudb.connect({str(path)!r}, read_only=True)\n" + "result = conn.execute('MATCH (p:person) RETURN p.uid AS uid')\n" + "print('before', 'numpy' in sys.modules)\n" + "arrays = result.fetchnumpy()\n" + "print('after', 'numpy' in sys.modules, len(arrays['uid']))\n" + ) + assert done.stdout.splitlines() == ["before False", "after True 3"] + + def test_pyarrow_arrives_when_a_result_is_asked_for_its_columns(tmp_path: Path) -> None: pytest.importorskip("pyarrow") path = tmp_path / "columns.zu1" diff --git a/tests/test_numpy.py b/tests/test_numpy.py new file mode 100644 index 0000000..402ad2d --- /dev/null +++ b/tests/test_numpy.py @@ -0,0 +1,175 @@ +"""A result as numpy arrays. + +The third way out of a result, and the one with the fewest dependencies +under it: numpy and nothing else. What these check is the mapping, which +is where a columnar export goes wrong quietly, and the two claims worth +making about it, which are that the buffer is not copied on the way and +that a null does not become a number. +""" + +from __future__ import annotations + +import datetime + +import pytest +import zudb + +np = pytest.importorskip("numpy") + + +def test_every_column_comes_back_under_its_own_name(social: zudb.Connection) -> None: + result = social.execute("MATCH (p:person) RETURN p.uid AS uid, p.name AS name ORDER BY p.uid") + arrays = result.fetchnumpy() + assert list(arrays) == ["uid", "name"], "the order the statement projected them in" + assert arrays["uid"].tolist() == [10, 20, 30] + assert arrays["name"].tolist() == ["ada", "grace", "kay"] + + +@pytest.mark.parametrize( + "statement,params,dtype,answer", + [ + ("RETURN 1 AS v", {}, "int64", 1), + ("RETURN 1.5 AS v", {}, "float64", 1.5), + ("RETURN true AS v", {}, "bool", True), + ("RETURN 'ada' AS v", {}, "object", "ada"), + ("RETURN $v AS v", {"v": datetime.date(2020, 1, 2)}, "datetime64[D]", "2020-01-02"), + ( + "RETURN $v AS v", + {"v": datetime.datetime(2020, 1, 2, 3, 4, 5)}, + "datetime64[ns]", + "2020-01-02T03:04:05.000000000", + ), + ("RETURN $v AS v", {"v": datetime.time(1, 2, 3)}, "timedelta64[ns]", 3723000000000), + ( + "RETURN $v AS v", + {"v": datetime.timedelta(hours=1, microseconds=5)}, + "timedelta64[ns]", + 3600000005000, + ), + ("RETURN DURATION 'P1Y2M' AS v", {}, "timedelta64[M]", 14), + ("RETURN null AS v", {}, "object", None), + ], +) +def test_a_column_becomes_the_numpy_type_that_holds_it( + empty: zudb.Connection, statement: str, params: dict, dtype: str, answer: object +) -> None: + array = empty.execute(statement, params).fetchnumpy()["v"] + assert array.dtype == np.dtype(dtype) + assert array[0] == np.array([answer], dtype=dtype)[0] + + +def test_a_time_of_day_is_nanoseconds_since_midnight(empty: zudb.Connection) -> None: + """numpy has no clock reading, and this is the reading and not a stand-in.""" + array = empty.execute("RETURN $v AS v", {"v": datetime.time(1, 2, 3, 500000)}).fetchnumpy()["v"] + assert array[0] == np.timedelta64(3723500000000, "ns") + assert array[0].astype("timedelta64[s]") == np.timedelta64(3723, "s") + + +def test_a_datetime_with_an_offset_comes_back_as_the_instant_in_utc( + empty: zudb.Connection, +) -> None: + """numpy has nowhere to keep the offset, and the instant is what the buffer holds.""" + zone = datetime.timezone(datetime.timedelta(hours=5, minutes=30)) + when = datetime.datetime(2020, 1, 2, 3, 4, 5, tzinfo=zone) + array = empty.execute("RETURN $v AS v", {"v": when}).fetchnumpy()["v"] + assert array[0] == np.datetime64("2020-01-01T21:34:05.000000000") + + +def test_a_time_with_an_offset_is_refused_rather_than_moved(empty: zudb.Connection) -> None: + zone = datetime.timezone(datetime.timedelta(hours=2)) + with pytest.raises(TypeError, match="time with an offset"): + empty.execute("RETURN $v AS v", {"v": datetime.time(1, 2, 3, tzinfo=zone)}).fetchnumpy() + + +@pytest.mark.parametrize( + "statement,dtype,answer", + [ + ("UNWIND [1, null, 3] AS v RETURN v", "int64", [1, 3]), + ("UNWIND [1.5, null, 3.5] AS v RETURN v", "float64", [1.5, 3.5]), + ("UNWIND [true, null, false] AS v RETURN v", "bool", [True, False]), + ], +) +def test_a_column_with_a_null_in_it_is_masked( + empty: zudb.Connection, statement: str, dtype: str, answer: list +) -> None: + """numpy has no missing integer, so the mask is where the null goes.""" + array = empty.execute(statement).fetchnumpy()["v"] + assert isinstance(array, np.ma.MaskedArray) + assert array.dtype == np.dtype(dtype) + assert array.mask.tolist() == [False, True, False] + assert array.compressed().tolist() == answer + + +def test_a_column_with_nothing_missing_is_a_plain_array(empty: zudb.Connection) -> None: + """A mask nothing is masked by is a second buffer nobody asked for.""" + array = empty.execute("UNWIND [1, 2, 3] AS v RETURN v").fetchnumpy()["v"] + assert not isinstance(array, np.ma.MaskedArray) + assert array.tolist() == [1, 2, 3] + + +@pytest.mark.parametrize( + "statement,answer", + [ + ("UNWIND ['a', null, 'ccc'] AS v RETURN v", ["a", None, "ccc"]), + ("UNWIND [[1, 2], null] AS v RETURN v", [[1, 2], None]), + ("UNWIND [{a: 1}, null] AS v RETURN v", [{"a": 1}, None]), + ], +) +def test_an_object_column_carries_the_null_in_the_cell( + empty: zudb.Connection, statement: str, answer: list +) -> None: + """An object array has somewhere to put `None`, so a mask would say it twice.""" + array = empty.execute(statement).fetchnumpy()["v"] + assert array.dtype == np.dtype("object") + assert not isinstance(array, np.ma.MaskedArray) + assert array.tolist() == answer + + +def test_nodes_and_rels_come_back_as_the_objects_the_rows_hold(loaded: zudb.Connection) -> None: + statement = "MATCH (p:person)-[k:knows]->(q:person) RETURN p AS p, k AS k" + arrays = loaded.execute(statement).fetchnumpy() + assert arrays["p"].dtype == np.dtype("object") + assert all(isinstance(node, zudb.Node) for node in arrays["p"]) + assert all(isinstance(rel, zudb.Rel) for rel in arrays["k"]) + + +def test_the_buffer_is_moved_into_numpy_and_not_copied(empty: zudb.Connection) -> None: + """The claim the whole path is for: the array is the engine's buffer. + + `owndata` false with a base object is numpy's own way of saying the + memory came from somewhere else and is being kept alive by whoever + it came from, which here is the extension holding the `Vec`. + """ + array = empty.execute("UNWIND [1, 2, 3] AS v RETURN v").fetchnumpy()["v"] + assert not array.flags.owndata + assert array.base is not None + # And it is the caller's to write in: the result kept no second + # reference to it, so there is nothing for a write to corrupt. + assert array.flags.writeable + array[0] = 99 + assert array.tolist() == [99, 2, 3] + + +def test_a_result_with_no_rows_gives_empty_arrays_of_the_right_type( + empty: zudb.Connection, +) -> None: + arrays = empty.execute("UNWIND [] AS v RETURN v").fetchnumpy() + assert list(arrays) == ["v"] + assert len(arrays["v"]) == 0 + + +def test_a_statement_that_writes_gives_no_columns(empty: zudb.Connection) -> None: + assert empty.execute("INSERT (p:person {uid: 1, name: 'ada'})").fetchnumpy() == {} + + +def test_two_columns_of_the_same_name_are_refused(empty: zudb.Connection) -> None: + """A dict holds one of each name, and dropping the other one quietly is worse.""" + with pytest.raises(ValueError, match="two columns called 'v'"): + empty.execute("RETURN 1 AS v, 2 AS v").fetchnumpy() + + +def test_reading_the_columns_does_not_move_the_cursor(social: zudb.Connection) -> None: + """Like every other way of reading a result whole.""" + result = social.execute("MATCH (p:person) RETURN p.uid AS uid ORDER BY p.uid") + assert len(result.fetchnumpy()["uid"]) == 3 + assert result.fetchone() == (10,) diff --git a/tools/smoke.py b/tools/smoke.py index 4e16af9..20182ac 100644 --- a/tools/smoke.py +++ b/tools/smoke.py @@ -153,7 +153,7 @@ def nothing_else_installed() -> None: """ import zudb - for module in ("pandas", "polars", "pyarrow"): + for module in ("numpy", "pandas", "polars", "pyarrow"): assert module not in sys.modules, f"{module} was imported by importing zudb" with tempfile.TemporaryDirectory() as where: @@ -166,6 +166,16 @@ def nothing_else_installed() -> None: else: raise AssertionError("to_pandas worked without pandas") + # The extension links numpy's C API, which is the one of + # these that could plausibly be needed at import time and + # is not. + try: + rows.fetchnumpy() + except ImportError as refusal: + assert "zudb[numpy]" in str(refusal), refusal + else: + raise AssertionError("fetchnumpy worked without numpy") + def main() -> int: imported_from_an_install()