diff --git a/Cargo.lock b/Cargo.lock index 519ca82..92d1c36 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1709,7 +1709,7 @@ dependencies = [ [[package]] name = "zu" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "zu-common", "zu-encoding", @@ -1722,10 +1722,20 @@ dependencies = [ "zu-zu1", ] +[[package]] +name = "zu-arrow" +version = "0.0.1" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" +dependencies = [ + "arrow", + "zu-common", + "zu-query", +] + [[package]] name = "zu-common" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "thiserror", ] @@ -1733,7 +1743,7 @@ dependencies = [ [[package]] name = "zu-encoding" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "ruzstd", "zu-common", @@ -1742,7 +1752,7 @@ dependencies = [ [[package]] name = "zu-exec" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "zu-common", "zu-query", @@ -1752,7 +1762,7 @@ dependencies = [ [[package]] name = "zu-query" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "crossbeam-deque", "zu-common", @@ -1763,7 +1773,7 @@ dependencies = [ [[package]] name = "zu-s3" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "crc32c", "object_store", @@ -1774,7 +1784,7 @@ dependencies = [ [[package]] name = "zu-sqlite" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "rusqlite", "zu-common", @@ -1784,7 +1794,7 @@ dependencies = [ [[package]] name = "zu-storage" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "zu-common", "zu-encoding", @@ -1793,7 +1803,7 @@ dependencies = [ [[package]] name = "zu-vector" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "zu-common", ] @@ -1801,7 +1811,7 @@ dependencies = [ [[package]] name = "zu-zu1" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=130f67db924bcd0f766ee814b0da2edae32150d4#130f67db924bcd0f766ee814b0da2edae32150d4" +source = "git+https://github.com/tamnd/zu?rev=0698a4eccd31670f0a875b6d097d7753440a51b3#0698a4eccd31670f0a875b6d097d7753440a51b3" dependencies = [ "crc32c", "loom", @@ -1817,5 +1827,6 @@ dependencies = [ "numpy", "pyo3", "zu", + "zu-arrow", "zu-common", ] diff --git a/Cargo.toml b/Cargo.toml index 1c9021c..afc30a5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/README.md b/README.md index 01665d2..d32f5a0 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/conformance/cases.py b/conformance/cases.py index 45ae31e..ba1321d 100644 --- a/conformance/cases.py +++ b/conformance/cases.py @@ -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 = [] diff --git a/conformance/reader.py b/conformance/reader.py index 70ff294..0dd0e7b 100644 --- a/conformance/reader.py +++ b/conformance/reader.py @@ -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: diff --git a/python/zudb/_zudb.pyi b/python/zudb/_zudb.pyi index 35c3d1d..d0dca76 100644 --- a/python/zudb/_zudb.pyi +++ b/python/zudb/_zudb.pyi @@ -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: diff --git a/src/columns.rs b/src/columns.rs index 7cbfefe..1d8698c 100644 --- a/src/columns.rs +++ b/src/columns.rs @@ -15,80 +15,55 @@ //! version of pyarrow it has to agree with, because the interface is a //! C struct and not a Python API. //! -//! The columns themselves are not built here. `zudb::query::column` -//! reads a result down its columns in the engine, in two passes over -//! the rows, and hands back one owned buffer per column in the layout -//! Arrow already uses: values end to end, a validity bitmap that is -//! absent when nothing is null, strings as bytes and offsets. This -//! module takes those buffers and puts an Arrow array around them, -//! which for integers, floats, booleans, strings, dates, times, -//! datetimes and durations is a move and not a copy. `docs/clients/duckdb.md` -//! in the engine tree is why: this file used to walk the whole result -//! once per column to infer a type and once per column per batch to -//! gather pointers, and that transpose was the twenty. +//! ## The translation is not written here //! -//! What is left to build by hand is what no buffer covers: nodes, rels, -//! paths, lists and records, which arrive as borrowed values and become -//! structs and lists the way they always did. They are also the columns -//! nobody exports a million of. +//! It used to be, seven hundred lines of it, and the JavaScript client +//! had its own seven hundred that had to agree with them. `zu-arrow` in +//! the engine tree is the one answer about what a zu column becomes in +//! Arrow, and both clients export through it, so a year-month duration +//! is a month interval in both and a node names its table in both. The +//! engine is also where the buffers are: `zudb::query::column` reads a +//! result down its columns in two passes and hands back one owned +//! buffer per column in the layout Arrow already uses, and putting an +//! array around one of those is a move rather than a copy. //! -//! A column has one type, which the engine decides and this module only -//! translates. Two refusals stay here, because they are Arrow's facts -//! and not the engine's: a time with an offset has no Arrow type, and -//! neither has a handle to a graph or a binding table. +//! What is left here is the three things a shared crate cannot know: +//! which Python exception each kind of refusal is, where the table +//! names live in this client, and how many rows a caller wanted in a +//! batch. -use std::collections::HashMap; -use std::sync::Arc; - -use arrow::array::{ - Array, ArrayRef, BooleanArray, Date32Array, DurationNanosecondArray, Float64Array, Int64Array, - IntervalMonthDayNanoArray, LargeStringArray, ListArray, NullArray, StringArray, StructArray, - Time64NanosecondArray, TimestampNanosecondArray, UInt64Array, -}; -use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; -use arrow::datatypes::{ - DataType, Field, FieldRef, Fields, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, - TimeUnit, -}; -use arrow::error::ArrowError; -use arrow::ffi_stream::FFI_ArrowArrayStream; -use arrow::record_batch::{RecordBatch, RecordBatchOptions, RecordBatchReader}; use pyo3::prelude::*; -use zu_common::{DurationKind, Temporal}; -use zudb::query::column::{ColumnData, ColumnType, Offsets, Validity}; -use zudb::query::{QueryResult, Value}; use crate::value::Names; -/// How many rows go in one record batch. +/// What goes wrong here, with the GIL down and no way to raise yet. /// -/// A result is already in memory and the arrays are built whole, so a -/// batch is a view into them rather than a copy: the boundary exists -/// because readers expect one and because a working set that fits in -/// cache is faster to consume, not because anything is allocated at it. -const BATCH: usize = 65_536; +/// The engine's error, under the name this client has always called it, +/// because what it means has not changed: something in a column could +/// not be said in Arrow. +pub use zu_arrow::Error as Snag; -/// What goes wrong here, with the GIL down and no way to raise yet. +/// Turning one into the exception it is. /// -/// The two Python classes are the two mistakes: a value of the wrong -/// type in a column is a `TypeError`, and a value of the right type -/// that will not fit is a `ValueError`. Arrow's own errors are neither, -/// and are internal until one of them turns out to be reachable. -pub enum Snag { - Type(String), - Value(String), - Arrow(ArrowError), +/// An extension trait rather than a method, because [`Snag`] belongs to +/// the engine now and only this client knows that a value of the wrong +/// type is a `TypeError` here. +pub trait Raise { + /// The exception this is, once there is a GIL to raise it with. + fn raise(self, py: Python<'_>) -> PyErr; } -impl Snag { - /// The exception this is, once there is a GIL to raise it with. - pub fn raise(self, _py: Python<'_>) -> PyErr { +impl Raise for Snag { + fn raise(self, _py: Python<'_>) -> PyErr { match self { + // The two Python classes are the two mistakes: a value of + // the wrong type in a column is a `TypeError`, and a value + // of the right type that will not fit is a `ValueError`. Snag::Type(detail) => pyo3::exceptions::PyTypeError::new_err(detail), Snag::Value(detail) => pyo3::exceptions::PyValueError::new_err(detail), // Nothing here is meant to be reachable: the types are // decided before a buffer is filled, so an Arrow error is - // this module getting it wrong rather than the caller. + // the translation getting it wrong rather than the caller. Snag::Arrow(err) => pyo3::exceptions::PyRuntimeError::new_err(format!( "arrow could not build the result: {err}" )), @@ -96,617 +71,38 @@ impl Snag { } } -impl From for Snag { - fn from(err: ArrowError) -> Snag { - Snag::Arrow(err) - } -} - -/// A buffer that does not match the type the engine decided for it, -/// which is this module reading its own input wrong. -fn mismatch(name: &str, ty: &ColumnType) -> Snag { - Snag::Arrow(ArrowError::SchemaError(format!( - "column '{name}' came back as {} in a buffer that does not hold one", - ty.name() - ))) -} - -/// The refusal for a type Arrow has nowhere to put. +/// The names this client already took off the catalog, offered to the +/// translation in the shape it asks for. /// -/// Two of them, and both are Arrow's facts rather than the engine's, -/// which is why they live in the client and not in `columnar()`. -fn unsupported(name: &str, ty: &ColumnType) -> Snag { - match ty { - // Arrow has a time and a timestamp and nothing in between: - // there is no time-with-offset type to put this in, and - // dropping the offset would move the value. - ColumnType::ZonedTime { .. } => Snag::Type(format!( - "column '{name}' holds a time with an offset, which Arrow has no type for" - )), - // GV60 and GV61. A handle is a reference, and a column of - // references is a column of nothing a frame can hold: the graph - // is in the file and the binding table is behind the handle. A - // caller who wants one in a frame reads the rows, where it - // arrives as the string that names it, or projects the columns - // of the table instead of the table. - ColumnType::Graph | ColumnType::BindingTable => Snag::Type(format!( - "column '{name}' holds a reference to a graph or a binding table, which Arrow has no type for" - )), - _ => mismatch(name, ty), +/// Borrowed rather than cloned, because a column of a hundred million +/// nodes is a hundred million lookups. The inherent `node` and `rel` on +/// [`Names`] answer with the `#id` fallback and stay where they are; +/// these two are the raw question, which is what the translation wants +/// so it can decide the fallback itself. +impl zu_arrow::Tables for Names { + fn node(&self, id: u32) -> Option<&str> { + self.node_name(id) } -} - -/// Every field is nullable, here and in the nested types, because a -/// null row of a struct column is a null in each of its children and -/// there is no other place to put it. -fn field(name: &str, data_type: DataType) -> FieldRef { - Arc::new(Field::new(name, data_type, true)) -} - -fn item(data_type: DataType) -> FieldRef { - field("item", data_type) -} -fn node_fields() -> Fields { - Fields::from(vec![ - field("table", DataType::Utf8), - field("offset", DataType::UInt64), - ]) -} - -fn rel_fields() -> Fields { - Fields::from(vec![ - field("table", DataType::Utf8), - field("src", DataType::UInt64), - field("dst", DataType::UInt64), - field("ord", DataType::UInt64), - ]) -} - -fn path_fields() -> Fields { - Fields::from(vec![ - field( - "nodes", - DataType::List(item(DataType::Struct(node_fields()))), - ), - field("rels", DataType::List(item(DataType::Struct(rel_fields())))), - ]) -} - -/// An offset in minutes as the name Arrow keeps a timezone under. -/// -/// A fixed offset rather than a region, because a fixed offset is what -/// the value carries: the engine stores when a zoned datetime happened -/// and how far from UTC it was written, and no amount of arithmetic -/// recovers `Europe/Paris` from `+01:00`. -fn zone(offset: i16) -> String { - let sign = if offset < 0 { '-' } else { '+' }; - let minutes = offset.unsigned_abs(); - format!("{sign}{:02}:{:02}", minutes / 60, minutes % 60) + fn rel(&self, id: u32) -> Option<&str> { + self.rel_name(id) + } } -/// The Arrow type a column type becomes, and the two places where the -/// answer is that it does not become one. -/// -/// The column name rides along because a refusal without it sends -/// somebody to read a schema by hand, and because a nested refusal is -/// still about the column it is nested in. -fn data_type(name: &str, ty: &ColumnType) -> Result { - Ok(match ty { - ColumnType::Null => DataType::Null, - ColumnType::Bool => DataType::Boolean, - ColumnType::Int => DataType::Int64, - ColumnType::Float => DataType::Float64, - ColumnType::Str => DataType::Utf8, - ColumnType::Date => DataType::Date32, - ColumnType::LocalTime => DataType::Time64(TimeUnit::Nanosecond), - ColumnType::LocalDatetime => DataType::Timestamp(TimeUnit::Nanosecond, None), - ColumnType::ZonedDatetime { offset } => { - DataType::Timestamp(TimeUnit::Nanosecond, Some(zone(*offset).into())) - } - // Arrow has a year-month interval, which is exactly what this - // is, and pyarrow cannot build a Python array of one: its type - // id has no class behind it, so reading such a column raises - // `KeyError: 21`. Month-day-nano is the interval every reader - // implements, and a year-month duration is one with no days and - // no nanoseconds in it. - ColumnType::YearMonth => DataType::Interval(IntervalUnit::MonthDayNano), - ColumnType::DayTime => DataType::Duration(TimeUnit::Nanosecond), - ColumnType::Node => DataType::Struct(node_fields()), - ColumnType::Rel => DataType::Struct(rel_fields()), - ColumnType::Path => DataType::Struct(path_fields()), - ColumnType::List(of) => DataType::List(item(data_type(name, of)?)), - ColumnType::Record(fields) => DataType::Struct( - fields - .iter() - .map(|(held, ty)| Ok(field(held, data_type(name, ty)?))) - .collect::>()?, - ), - ColumnType::ZonedTime { .. } | ColumnType::Graph | ColumnType::BindingTable => { - return Err(unsupported(name, ty)); - } - }) -} +/// How many rows go in one record batch, when a caller has no opinion. +pub const BATCH: usize = zu_arrow::BATCH; /// The stream a result exports, batches and schema and all. /// /// One array per column, built once out of the engine's buffers, and /// batches that are slices of them. The arrays are built eagerly /// because the refusals have to happen while there is still a caller to -/// raise them at; the batches are not, so `record_batches` no longer -/// builds a second copy of the table before it hands back a reader. -pub fn stream(result: &QueryResult, names: &Names) -> Result { - let columns = result - .columnar() - .map_err(|mixed| Snag::Type(mixed.to_string()))?; - let rows = columns.rows; - - let mut fields = Vec::with_capacity(columns.len()); - let mut arrays = Vec::with_capacity(columns.len()); - for held in columns.columns { - let array = column( - held.name, - &held.ty, - held.data, - held.validity, - held.len, - names, - )?; - fields.push(field(held.name, array.data_type().clone())); - arrays.push(array); - } - - let schema = Arc::new(Schema::new(Fields::from(fields))); - Ok(FFI_ArrowArrayStream::new(Box::new(Slices { - schema, - arrays, - rows, - at: 0, - given: 0, - }))) -} - -/// The batches, cut out of the finished arrays as they are asked for. -/// -/// A result with no rows still has a schema, and a reader that gets no -/// batch at all cannot tell what the columns were, so an empty result -/// gives one empty batch and then stops. -struct Slices { - schema: SchemaRef, - arrays: Vec, - rows: usize, - at: usize, - given: usize, -} - -impl Iterator for Slices { - type Item = Result; - - fn next(&mut self) -> Option> { - if self.at >= self.rows && self.given > 0 { - return None; - } - let take = BATCH.min(self.rows - self.at); - let columns: Vec = self - .arrays - .iter() - .map(|array| array.slice(self.at, take)) - .collect(); - self.at += take; - self.given += 1; - // The row count goes in by hand because a result with no - // columns still has rows, and a batch of no columns cannot say - // how many any other way. - Some(RecordBatch::try_new_with_options( - self.schema.clone(), - columns, - &RecordBatchOptions::new().with_row_count(Some(take)), - )) - } -} - -impl RecordBatchReader for Slices { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} - -/// The bitmap Arrow keeps beside a buffer, out of the one the engine -/// filled. Absent means every row has a value, in both layouts. -fn nulls(validity: Option) -> Option { - validity - .map(|held| NullBuffer::new(BooleanBuffer::new(Buffer::from_vec(held.bits), 0, held.len))) -} - -/// One whole column as an Arrow array. -/// -/// Every flat arm here moves a `Vec` into an Arrow buffer and allocates -/// nothing: the engine filled it in the layout Arrow reads, and the -/// only work left is putting a type and a bitmap around it. The two -/// exceptions are year-month intervals, which are 96 bits in Arrow and -/// 64 in the engine, and the complex types, which have no buffer. -fn column( - name: &str, - ty: &ColumnType, - data: ColumnData<'_>, - validity: Option, - len: usize, +/// raise them at; the batches are not, so a reader that stops early +/// stops paying. +pub fn stream( + result: &zudb::query::QueryResult, names: &Names, -) -> Result { - let valid = nulls(validity); - Ok(match data { - ColumnData::Null => Arc::new(NullArray::new(len)), - ColumnData::Bool { bits } => Arc::new(BooleanArray::new( - BooleanBuffer::new(Buffer::from_vec(bits), 0, len), - valid, - )), - ColumnData::Int(values) => Arc::new(Int64Array::new(ScalarBuffer::from(values), valid)), - ColumnData::Float(values) => Arc::new(Float64Array::new(ScalarBuffer::from(values), valid)), - ColumnData::Str(held) => match held.offsets { - Offsets::I32(offsets) => Arc::new(StringArray::try_new( - OffsetBuffer::new(ScalarBuffer::from(offsets)), - Buffer::from_vec(held.bytes), - valid, - )?), - // Past two gigabytes of text in one column, which is where - // a 32 bit offset stops addressing the bytes. Arrow's own - // answer is the wider type and every reader has it. - Offsets::I64(offsets) => Arc::new(LargeStringArray::try_new( - OffsetBuffer::new(ScalarBuffer::from(offsets)), - Buffer::from_vec(held.bytes), - valid, - )?), - }, - ColumnData::Days(values) => Arc::new(Date32Array::new(ScalarBuffer::from(values), valid)), - ColumnData::Nanos(values) => { - let values = ScalarBuffer::from(values); - match ty { - ColumnType::LocalTime => Arc::new(Time64NanosecondArray::new(values, valid)), - ColumnType::LocalDatetime => Arc::new(TimestampNanosecondArray::new(values, valid)), - ColumnType::ZonedDatetime { offset } => Arc::new( - TimestampNanosecondArray::new(values, valid).with_timezone(zone(*offset)), - ), - ColumnType::DayTime => Arc::new(DurationNanosecondArray::new(values, valid)), - // A time with an offset fills a nanosecond buffer like - // any other time, and this is where it stops. - _ => return Err(unsupported(name, ty)), - } - } - ColumnData::Months(counts) => Arc::new(IntervalMonthDayNanoArray::new( - ScalarBuffer::from(months(name, &counts)?), - valid, - )), - // The types with no buffer: nodes, rels, paths, lists, records, - // and the two handles, which reach here as values and are - // refused there. - ColumnData::Complex(values) => build(name, ty, &values, names)?, - }) -} - -/// Month counts as the interval Arrow carries them in. -/// -/// Arrow counts the months of an interval in 32 bits and the engine -/// counts them in 64, so the far end of the range has nowhere to go. -/// Refusing it is the only honest answer; wrapping would move the value -/// by centuries. -fn months(name: &str, counts: &[i64]) -> Result, Snag> { - let mut months = Vec::with_capacity(counts.len()); - for (row, count) in counts.iter().enumerate() { - let count = i32::try_from(*count).map_err(|_| { - Snag::Value(format!( - "the duration at row {row} of column '{name}' is {count} months, which is more than an Arrow interval holds" - )) - })?; - months.push(IntervalMonthDayNano::new(count, 0, 0)); - } - Ok(months) -} - -/// One column's array, walked out of the values in it. -/// -/// This is the slow path and it is where the complex types live: the -/// top level reaches it only for nodes, rels, paths, lists and records, -/// and everything below the top level reaches it always, because a list -/// item and a record field are values wherever they sit. -fn build(name: &str, ty: &ColumnType, values: &[&Value], names: &Names) -> Result { - Ok(match ty { - ColumnType::Null => Arc::new(NullArray::new(values.len())), - ColumnType::Bool => Arc::new( - values - .iter() - .map(|value| match value { - Value::Bool(b) => Some(*b), - _ => None, - }) - .collect::(), - ), - ColumnType::Int => Arc::new( - values - .iter() - .map(|value| match value { - Value::Int(n) => Some(*n), - _ => None, - }) - .collect::(), - ), - ColumnType::Float => Arc::new( - values - .iter() - .map(|value| match value { - Value::Float(f) => Some(*f), - // Widened where the column holds both, which is - // the only place an integer reaches a float column. - Value::Int(n) => Some(*n as f64), - _ => None, - }) - .collect::(), - ), - ColumnType::Str => Arc::new( - values - .iter() - .map(|value| match value { - Value::Str(s) => Some(s.as_str()), - _ => None, - }) - .collect::(), - ), - ColumnType::Date => Arc::new( - temporals(values) - .map(|temporal| match temporal { - Some(Temporal::Date(days)) => Some(*days), - _ => None, - }) - .collect::(), - ), - ColumnType::LocalTime => Arc::new( - temporals(values) - .map(|temporal| match temporal { - Some(Temporal::LocalTime(nanos)) => Some(*nanos), - _ => None, - }) - .collect::(), - ), - ColumnType::LocalDatetime => Arc::new( - temporals(values) - .map(|temporal| match temporal { - Some(Temporal::LocalDatetime(nanos)) => Some(*nanos), - _ => None, - }) - .collect::(), - ), - ColumnType::ZonedDatetime { offset } => Arc::new( - temporals(values) - .map(|temporal| match temporal { - Some(Temporal::ZonedDatetime { nanos, .. }) => Some(*nanos), - _ => None, - }) - .collect::() - .with_timezone(zone(*offset)), - ), - ColumnType::YearMonth => { - let mut counts = Vec::with_capacity(values.len()); - let mut valid = Vec::with_capacity(values.len()); - for temporal in temporals(values) { - match temporal { - Some(Temporal::Duration(DurationKind::YearMonth, count)) => { - counts.push(*count); - valid.push(true); - } - _ => { - counts.push(0); - valid.push(false); - } - } - } - Arc::new(IntervalMonthDayNanoArray::new( - ScalarBuffer::from(months(name, &counts)?), - Some(NullBuffer::from(valid)), - )) - } - ColumnType::DayTime => Arc::new( - temporals(values) - .map(|temporal| match temporal { - Some(Temporal::Duration(DurationKind::DayTime, nanos)) => Some(*nanos), - _ => None, - }) - .collect::(), - ), - ColumnType::Node => nodes(values, names)?, - ColumnType::Rel => rels(values, names)?, - ColumnType::Path => paths(name, values, names)?, - ColumnType::List(of) => { - let mut offsets = Vec::with_capacity(values.len() + 1); - let mut flat: Vec<&Value> = Vec::new(); - let mut valid = Vec::with_capacity(values.len()); - offsets.push(0i32); - for value in values { - if let Value::List(items) = value { - flat.extend(items.iter()); - valid.push(true); - } else { - valid.push(false); - } - offsets.push(flat.len() as i32); - } - Arc::new(ListArray::try_new( - item(data_type(name, of)?), - OffsetBuffer::new(offsets.into()), - build(name, of, &flat, names)?, - Some(NullBuffer::from(valid)), - )?) - } - ColumnType::Record(fields) => { - let mut children: Vec = Vec::with_capacity(fields.len()); - for (at, (_, ty)) in fields.iter().enumerate() { - let column: Vec<&Value> = values - .iter() - .map(|value| match value { - Value::Record(held) => &held[at].1, - _ => &Value::Null, - }) - .collect(); - children.push(build(name, ty, &column, names)?); - } - Arc::new(StructArray::try_new( - match data_type(name, ty)? { - DataType::Struct(fields) => fields, - _ => return Err(mismatch(name, ty)), - }, - children, - Some(present(values)), - )?) - } - ColumnType::ZonedTime { .. } | ColumnType::Graph | ColumnType::BindingTable => { - return Err(unsupported(name, ty)); - } - }) -} - -/// The temporal each value holds, or `None` for a value that is not one -/// and for a null. -fn temporals<'a>(values: &'a [&'a Value]) -> impl Iterator> { - values.iter().map(|value| match value { - Value::Temporal(temporal) => Some(temporal), - _ => None, - }) -} - -/// Which rows of a struct column are there at all. -fn present(values: &[&Value]) -> NullBuffer { - NullBuffer::from( - values - .iter() - .map(|value| !matches!(value, Value::Null)) - .collect::>(), - ) -} - -fn nodes(values: &[&Value], names: &Names) -> Result { - let table = tables( - values, - |value| match value { - Value::Node { table, .. } => Some(*table), - _ => None, - }, - |id| names.node_name(id), - ); - let offset: UInt64Array = values - .iter() - .map(|value| match value { - Value::Node { offset, .. } => Some(*offset), - _ => None, - }) - .collect(); - Ok(Arc::new(StructArray::try_new( - node_fields(), - vec![Arc::new(table), Arc::new(offset)], - Some(present(values)), - )?)) -} - -fn rels(values: &[&Value], names: &Names) -> Result { - let table = tables( - values, - |value| match value { - Value::Rel { table, .. } => Some(*table), - _ => None, - }, - |id| names.rel_name(id), - ); - let end = |pick: fn(&Value) -> Option| -> UInt64Array { - values.iter().map(|value| pick(value)).collect() - }; - Ok(Arc::new(StructArray::try_new( - rel_fields(), - vec![ - Arc::new(table), - Arc::new(end(|value| match value { - Value::Rel { src, .. } => Some(*src), - _ => None, - })), - Arc::new(end(|value| match value { - Value::Rel { dst, .. } => Some(*dst), - _ => None, - })), - Arc::new(end(|value| match value { - Value::Rel { ord, .. } => Some(*ord), - _ => None, - })), - ], - Some(present(values)), - )?)) -} - -/// The table name of every row, borrowed rather than copied. -/// -/// The catalog owns the names and a column holds as many rows as the -/// result does, so the names go in by reference and the only string -/// built here is the stand-in for a table the catalog no longer has, -/// which is one per missing table rather than one per row. -fn tables<'a>( - values: &[&Value], - id_of: impl Fn(&Value) -> Option, - name_of: impl Fn(u32) -> Option<&'a str>, -) -> StringArray { - let mut gone: HashMap = HashMap::new(); - for value in values { - if let Some(id) = id_of(value) - && name_of(id).is_none() - { - gone.entry(id).or_insert_with(|| format!("#{id}")); - } - } - values - .iter() - .map(|value| id_of(value).map(|id| name_of(id).unwrap_or_else(|| gone[&id].as_str()))) - .collect() -} - -/// A path column, as the two lists a walk is. -/// -/// A path is nodes and edges alternating, and Arrow has no type for a -/// list whose elements alternate between two structs. Two lists say the -/// same thing without a union in the middle of it: the nodes in the -/// order the walk visits them, the edges in the order it crosses them, -/// and one more node than edge. -fn paths(name: &str, values: &[&Value], names: &Names) -> Result { - let mut node_offsets = vec![0i32]; - let mut rel_offsets = vec![0i32]; - let mut walked_nodes: Vec<&Value> = Vec::new(); - let mut walked_rels: Vec<&Value> = Vec::new(); - for (row, value) in values.iter().enumerate() { - match value { - Value::Path(elements) => { - walked_nodes.extend(elements.iter().step_by(2)); - walked_rels.extend(elements.iter().skip(1).step_by(2)); - } - // Never in a result: the executor settles a chain into its - // edges before the rows leave the pipeline. - Value::Chain(_) => { - return Err(Snag::Type(format!( - "row {row} of column '{name}' is a path chain, which is internal to the executor" - ))); - } - _ => {} - } - node_offsets.push(walked_nodes.len() as i32); - rel_offsets.push(walked_rels.len() as i32); - } - let nodes = ListArray::try_new( - item(DataType::Struct(node_fields())), - OffsetBuffer::new(node_offsets.into()), - nodes(&walked_nodes, names)?, - Some(present(values)), - )?; - let rels = ListArray::try_new( - item(DataType::Struct(rel_fields())), - OffsetBuffer::new(rel_offsets.into()), - rels(&walked_rels, names)?, - Some(present(values)), - )?; - Ok(Arc::new(StructArray::try_new( - path_fields(), - vec![Arc::new(nodes), Arc::new(rels)], - Some(present(values)), - )?)) + rows: usize, +) -> Result { + zu_arrow::stream(result, names, rows) } diff --git a/src/conn.rs b/src/conn.rs index 90f2a92..424bfde 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -12,13 +12,14 @@ use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex, OnceLock}; +use arrow::ffi_stream::FFI_ArrowArrayStream; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use zudb::query::{QueryResult, Value}; use zudb::{Config, Database, Interrupt}; use crate::appender::Appender; -use crate::columns; +use crate::columns::{self, Raise}; use crate::error::{closed, programming, to_py_err}; use crate::html; use crate::interrupt; @@ -799,7 +800,7 @@ impl Result { // in a DataFrame is big enough that another thread should get // to run while it happens. let stream = py - .detach(|| columns::stream(&self.result, &self.names)) + .detach(|| columns::stream(&self.result, &self.names, columns::BATCH)) .map_err(|snag| snag.raise(py))?; PyCapsule::new_with_value(py, stream, STREAM) } @@ -863,16 +864,34 @@ impl Result { /// /// The same data as `to_arrow`, handed over in batches of sixty-five /// thousand rows instead of as one table, which is what a consumer - /// that writes as it reads wants. + /// that writes as it reads wants. `rows_per_batch` says otherwise + /// for a consumer that has a size in mind, a Parquet row group or + /// whatever its own downstream takes. /// /// A batch is a view of the column it comes from rather than a copy /// of it, and it is cut when the reader asks for it, so a consumer - /// that stops early stops paying. - fn record_batches<'py>(slf: PyRef<'py, Self>) -> PyResult> { + /// that stops early stops paying, and the size costs nothing either + /// way. + #[pyo3(signature = (rows_per_batch = None))] + fn record_batches<'py>( + slf: PyRef<'py, Self>, + rows_per_batch: Option, + ) -> PyResult> { let py = slf.py(); - needed(py, "pyarrow", "arrow")? - .getattr("RecordBatchReader")? - .call_method1("from_stream", (slf,)) + let rows = match rows_per_batch { + Some(rows) if rows < 1 => return Err(programming(py, EMPTY_BATCH)), + Some(rows) => rows as usize, + None => columns::BATCH, + }; + let reader = needed(py, "pyarrow", "arrow")?.getattr("RecordBatchReader")?; + // Borrowed out of the reference before the GIL goes down, + // because what the other thread may touch is the result and not + // the Python object holding it. + let (result, names) = (&slf.result, &slf.names); + let stream = py + .detach(|| columns::stream(result, names, rows)) + .map_err(|snag| snag.raise(py))?; + reader.call_method1("from_stream", (Batches::holding(py, stream)?,)) } fn __len__(&self) -> usize { @@ -944,6 +963,67 @@ impl Result { } } +/// What a batch size of nothing gets told. +const EMPTY_BATCH: &str = "rows_per_batch has to be at least one, because a reader cutting a \ + result into batches of no rows never reaches the end of it"; + +/// A stream that has already been cut to a size, waiting to be taken. +/// +/// `__arrow_c_stream__` takes no arguments, because the PyCapsule +/// protocol has none to give it, so a caller who asked for a batch size +/// cannot be handed the result itself. This is what it is handed +/// instead: the same capsule, off a stream that was built with the size +/// the caller wanted, and nothing else. +/// +/// It is a one-shot. The protocol's rule is that the consumer owns the +/// stream from the moment it takes the capsule, so there is nothing +/// left here to hand to a second one. +#[pyclass(module = "zudb")] +pub struct Batches { + stream: Mutex>, +} + +impl Batches { + /// One holding a stream, as the Python object it has to be to reach + /// `RecordBatchReader.from_stream`. + fn holding(py: Python<'_>, stream: FFI_ArrowArrayStream) -> PyResult> { + Py::new( + py, + Batches { + stream: Mutex::new(Some(stream)), + }, + ) + } +} + +#[pymethods] +impl Batches { + /// The capsule, once. + /// + /// `requested_schema` is accepted and ignored for the same reason a + /// result ignores it: the columns have the types they have, and + /// casting them here would hide a conversion a caller can see. + #[pyo3(signature = (requested_schema = None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + let taken = self + .stream + .lock() + .map_err(|_| programming(py, TAKEN))? + .take() + .ok_or_else(|| programming(py, TAKEN))?; + PyCapsule::new_with_value(py, taken, STREAM) + } +} + +/// What a second reader of the same batches gets told. +const TAKEN: &str = "these batches have already been read, and a stream is handed over once: call \ + record_batches again for another reader"; + /// A module the caller has to have installed for this call, imported. /// /// A missing one is reported as the install that fixes it rather than diff --git a/src/frame.rs b/src/frame.rs index 33e3121..79694d2 100644 --- a/src/frame.rs +++ b/src/frame.rs @@ -45,7 +45,7 @@ use zu_common::{DurationKind, FloatBits, IntBits, LogicalType}; use zudb::{Column, Layout}; use crate::buffer::{self, type_name}; -use crate::columns::Snag; +use crate::columns::{Raise, Snag}; use crate::load; /// What a capsule holding an Arrow stream is called, which a consumer diff --git a/tests/test_arrow.py b/tests/test_arrow.py index 3399d3e..0195b07 100644 --- a/tests/test_arrow.py +++ b/tests/test_arrow.py @@ -201,9 +201,11 @@ def test_a_result_with_no_rows_still_has_its_columns(loaded: zudb.Connection) -> table = loaded.execute("MATCH (p:person) WHERE p.uid = 99 RETURN p.uid AS uid").to_arrow() assert table.num_rows == 0 assert table.column_names == ["uid"] - # Nothing said what the column holds, so it holds nothing, which - # Arrow has a type for. - assert table.schema.field("uid").type == pa.null() + # The type comes off the column and not off the rows, so a result + # that matched nothing still says what it would have held. That is + # what a consumer writing Parquet or appending to an existing table + # needs from a query that happened to find no rows. + assert table.schema.field("uid").type == pa.int64() @pytest.mark.parametrize( @@ -249,6 +251,28 @@ def test_the_batches_are_the_same_rows(tmp_path: Path) -> None: assert sum(batch.num_rows for batch in batches) == rows +def test_a_caller_can_say_how_big_a_batch_is(loaded: zudb.Connection) -> None: + reader = loaded.execute("MATCH (p:person) RETURN p.uid AS uid").record_batches(2) + batches = list(reader) + assert [batch.num_rows for batch in batches] == [2, 1] + assert pa.Table.from_batches(batches).column("uid").to_pylist() == [10, 20, 30] + + +def test_a_batch_of_no_rows_is_refused(loaded: zudb.Connection) -> None: + with pytest.raises(zudb.ProgrammingError, match="at least one"): + loaded.execute("MATCH (p:person) RETURN p.uid AS uid").record_batches(0) + + +def test_the_batches_of_a_result_read_once(loaded: zudb.Connection) -> None: + result = loaded.execute("MATCH (p:person) RETURN p.uid AS uid") + first = result.record_batches() + assert sum(batch.num_rows for batch in first) == 3 + # The result is still there and still holds its rows, so asking it + # again is a second reader and not an error. + second = result.record_batches() + assert sum(batch.num_rows for batch in second) == 3 + + def test_a_result_reads_as_arrow_and_as_objects_and_says_the_same_thing( loaded: zudb.Connection, ) -> None: