From ac34ba94586e0cb64d937ad5a3f37c8fc90fe294 Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Wed, 19 Aug 2026 13:47:12 +0700 Subject: [PATCH] Hand Arrow the buffers the engine already filled docs/clients/duckdb.md measured to_arrow on a million rows at twenty times DuckDB's, and it traced all of it to one thing: a result is rows, and this file was the code that turned it back into columns. It walked the whole result once per column to settle a type, then once per column per batch to gather a Vec of pointers, then built an Arrow array by collecting an iterator of Options into it. Three passes over three million boxed values, and a copy at the end of them. The engine reads a result down its columns now. zudb::query::column does it in two passes over the rows in row order, and hands back one owned buffer per column in the layout Arrow already uses: values end to end with a null still occupying its cell, a validity bitmap that is absent when nothing is null, strings as bytes and offsets that stay narrow until the bytes pass what a 32 bit offset addresses. That is the half of the work this file used to do badly, done once, in the place that will eventually not have to do it at all. So what is left here is a translation. An integer column is Int64Array::new(ScalarBuffer::from(values), nulls) and the Vec moves; so do floats, dates, times, datetimes, day-time durations, the packed bits of a boolean column, and the bytes and offsets of a string one. Nothing is copied and nothing is walked. The two that still are: a year-month duration, because Arrow counts an interval's months in 32 bits and the engine counts them in 64, and the complex types, which have no buffer and arrive as borrowed values. Nodes, rels, paths, lists and records are built exactly as they were, and they are also the columns nobody exports a million of. Both refusals that were here stay here, because they are Arrow's facts rather than the engine's. A time with an offset has no Arrow type and dropping the offset would move the value. A handle to a graph or a binding table is a reference to something that is not in the frame. The mixed column message now comes from the engine, and says the same thing in the same shape, which is what the tests that read it were pinning. record_batches stops being worse than its name promised. It used to build every batch into a Vec and then hand back a reader over it, so a caller who asked for batches got the whole Arrow table built first and the rows still alive beside it. A batch is a slice of the finished column now, cut when the reader asks for it, which is a view and not an allocation, and a consumer that stops early stops paying. Measured with tools/versus_duckdb.py, which is in this commit so the numbers can be reproduced rather than believed. A million rows of an integer, a double and a five byte string, in a stored table on both sides, release build, fastest of five, DuckDB 1.5.5. execute and Arrow table 148 ms -> 73 ms 10.8x -> 4.9x slower execute and pandas 148 ms -> 79 ms 3.1x -> 1.5x slower The statement itself is 45 ms of each of those, so the export went from 103 ms to 26 ms. Of the 26, the engine's own transpose is 22, which is what the columnar bench in the engine tree times. The Arrow half is about four milliseconds now and there is not much left in it. The rest goes when the sink stops flattening its vectors into rows, and that is the second half of item 1 in the audit rather than anything this file can do. Two things the script found that are worth writing down. DuckDB's arrow() hands back a RecordBatchReader that has read nothing, so it times at nought and is not the same call as ours; to_arrow_table is, and that is what is compared. And register is nine times faster than DuckDB's for a frame of numbers and four times slower for one with a string column in it, because the string bytes get a UTF-8 validation pass on the way in and the numbers get nothing. Both rows are printed rather than the average of them. Six tests. A null in the middle of an integer, a float, a boolean and a string column, which is the validity bitmap and, for the boolean one, the only place two different bit packings sit next to each other. And a column with nothing missing carrying no bitmap at all, which is Arrow's convention and the engine's and is worth pinning because it is the case a reader gets to skip work for. --- Cargo.lock | 20 +- Cargo.toml | 4 +- src/columns.rs | 618 +++++++++++++++++++++-------------------- src/conn.rs | 4 + tests/test_arrow.py | 31 +++ tools/versus_duckdb.py | 194 +++++++++++++ 6 files changed, 564 insertions(+), 307 deletions(-) create mode 100644 tools/versus_duckdb.py diff --git a/Cargo.lock b/Cargo.lock index 6921392..96e44d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1647,7 +1647,7 @@ dependencies = [ [[package]] name = "zu" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "zu-common", "zu-encoding", @@ -1663,7 +1663,7 @@ dependencies = [ [[package]] name = "zu-common" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "thiserror", ] @@ -1671,7 +1671,7 @@ dependencies = [ [[package]] name = "zu-encoding" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "ruzstd", "zu-common", @@ -1680,7 +1680,7 @@ dependencies = [ [[package]] name = "zu-exec" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "zu-common", "zu-query", @@ -1690,7 +1690,7 @@ dependencies = [ [[package]] name = "zu-query" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "crossbeam-deque", "zu-common", @@ -1701,7 +1701,7 @@ dependencies = [ [[package]] name = "zu-s3" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "crc32c", "object_store", @@ -1712,7 +1712,7 @@ dependencies = [ [[package]] name = "zu-sqlite" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "rusqlite", "zu-common", @@ -1722,7 +1722,7 @@ dependencies = [ [[package]] name = "zu-storage" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "zu-common", "zu-encoding", @@ -1731,7 +1731,7 @@ dependencies = [ [[package]] name = "zu-vector" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "zu-common", ] @@ -1739,7 +1739,7 @@ dependencies = [ [[package]] name = "zu-zu1" version = "0.0.1" -source = "git+https://github.com/tamnd/zu?rev=6f38950b9f1c4b6cbca37f404ee660370a945c12#6f38950b9f1c4b6cbca37f404ee660370a945c12" +source = "git+https://github.com/tamnd/zu?rev=a1f310cf7be9bd422b6f4c8307b11eb97c2a23be#a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" dependencies = [ "crc32c", "loom", diff --git a/Cargo.toml b/Cargo.toml index f344e5d..0bf3d1a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,8 +18,8 @@ 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 = "6f38950b9f1c4b6cbca37f404ee660370a945c12" } -zu-common = { git = "https://github.com/tamnd/zu", rev = "6f38950b9f1c4b6cbca37f404ee660370a945c12" } +zudb = { package = "zu", git = "https://github.com/tamnd/zu", rev = "a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" } +zu-common = { git = "https://github.com/tamnd/zu", rev = "a1f310cf7be9bd422b6f4c8307b11eb97c2a23be" } # `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 diff --git a/src/columns.rs b/src/columns.rs index df71e1d..7cbfefe 100644 --- a/src/columns.rs +++ b/src/columns.rs @@ -15,41 +15,57 @@ //! version of pyarrow it has to agree with, because the interface is a //! C struct and not a Python API. //! -//! A column has one type, which the values decide: the first one that -//! is not null settles it and every value after it has to fit. Integers -//! widen to floats where a column holds both, since that is the one -//! mixture a projection produces by accident and the one no reader is -//! surprised by. Everything else that does not fit is refused, naming -//! the column and the row, because a column that quietly became strings -//! is worse than one that would not build. +//! 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. +//! +//! 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. +//! +//! 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. use std::collections::HashMap; use std::sync::Arc; use arrow::array::{ - ArrayRef, BooleanArray, Date32Array, DurationNanosecondArray, Float64Array, Int64Array, - IntervalMonthDayNanoArray, ListArray, NullArray, StringArray, StructArray, + Array, ArrayRef, BooleanArray, Date32Array, DurationNanosecondArray, Float64Array, Int64Array, + IntervalMonthDayNanoArray, LargeStringArray, ListArray, NullArray, StringArray, StructArray, Time64NanosecondArray, TimestampNanosecondArray, UInt64Array, }; -use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::buffer::{BooleanBuffer, Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ - DataType, Field, FieldRef, Fields, IntervalMonthDayNano, IntervalUnit, Schema, TimeUnit, + DataType, Field, FieldRef, Fields, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, + TimeUnit, }; use arrow::error::ArrowError; use arrow::ffi_stream::FFI_ArrowArrayStream; -use arrow::record_batch::{RecordBatch, RecordBatchIterator}; +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. /// -/// A result is already in memory, so this is not about streaming a -/// table too big to hold: it is about the copy. Batching keeps the -/// Arrow buffers a reader has to allocate down to a working set that -/// fits in cache, and it is what every other Arrow producer does. +/// 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; /// What goes wrong here, with the GIL down and no way to raise yet. @@ -86,92 +102,37 @@ impl From for Snag { } } -/// The type of one column, as this module thinks about it. -/// -/// Arrow's `DataType` is what it turns into, but not what it is -/// decided as: a node, a rel and a record all become structs, and -/// telling them apart afterwards by their field names would be reading -/// tea leaves. Deciding it once and carrying it is also what makes the -/// second pass, the one that fills the buffers, a match with no -/// re-inspection of the values in it. -#[derive(Clone, PartialEq)] -enum Kind { - /// Nothing but nulls, which Arrow has a type for. - Null, - Bool, - Int, - Float, - Str, - Date, - Time, - LocalDatetime, - /// A datetime with an offset, in minutes from UTC. The values are - /// instants, so the offset is how the column prints and not what it - /// holds; the first one in the column names the zone. - ZonedDatetime(i16), - YearMonth, - DayTime, - Node, - Rel, - Path, - List(Box), - Record(Vec<(String, Kind)>), +/// 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() + ))) } -impl Kind { - fn name(&self) -> String { - match self { - Kind::Null => "nulls".into(), - Kind::Bool => "booleans".into(), - Kind::Int => "integers".into(), - Kind::Float => "floats".into(), - Kind::Str => "strings".into(), - Kind::Date => "dates".into(), - Kind::Time => "times".into(), - Kind::LocalDatetime => "datetimes".into(), - Kind::ZonedDatetime(_) => "zoned datetimes".into(), - Kind::YearMonth => "year-month durations".into(), - Kind::DayTime => "day-time durations".into(), - Kind::Node => "nodes".into(), - Kind::Rel => "rels".into(), - Kind::Path => "paths".into(), - Kind::List(of) => format!("lists of {}", of.name()), - Kind::Record(_) => "records".into(), - } - } - - fn data_type(&self) -> DataType { - match self { - Kind::Null => DataType::Null, - Kind::Bool => DataType::Boolean, - Kind::Int => DataType::Int64, - Kind::Float => DataType::Float64, - Kind::Str => DataType::Utf8, - Kind::Date => DataType::Date32, - Kind::Time => DataType::Time64(TimeUnit::Nanosecond), - Kind::LocalDatetime => DataType::Timestamp(TimeUnit::Nanosecond, None), - Kind::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. - Kind::YearMonth => DataType::Interval(IntervalUnit::MonthDayNano), - Kind::DayTime => DataType::Duration(TimeUnit::Nanosecond), - Kind::Node => DataType::Struct(node_fields()), - Kind::Rel => DataType::Struct(rel_fields()), - Kind::Path => DataType::Struct(path_fields()), - Kind::List(of) => DataType::List(item(of.data_type())), - Kind::Record(fields) => DataType::Struct( - fields - .iter() - .map(|(name, kind)| Arc::new(Field::new(name, kind.data_type(), true))) - .collect(), - ), - } +/// The refusal for a type Arrow has nowhere to put. +/// +/// 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), } } @@ -224,177 +185,233 @@ fn zone(offset: i16) -> String { format!("{sign}{:02}:{:02}", minutes / 60, minutes % 60) } +/// 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)); + } + }) +} + /// The stream a result exports, batches and schema and all. /// -/// Built whole rather than lazily: the rows are already in memory, so a -/// reader that pulls one batch at a time would only be deferring a copy -/// it is going to ask for anyway, and building it here is what lets the -/// refusals happen while there is still a caller to raise them at. +/// 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 kinds = result - .columns - .iter() - .enumerate() - .map(|(at, name)| infer(name, result.rows.iter().map(|row| &row[at]))) - .collect::, Snag>>()?; - let schema = Arc::new(Schema::new( - result - .columns - .iter() - .zip(&kinds) - .map(|(name, kind)| field(name, kind.data_type())) - .collect::(), - )); - - let mut batches = Vec::new(); - let mut at = 0; - while at < result.rows.len() { - let rows = &result.rows[at..(at + BATCH).min(result.rows.len())]; - let columns = kinds - .iter() - .enumerate() - .map(|(ix, kind)| { - let values: Vec<&Value> = rows.iter().map(|row| &row[ix]).collect(); - build(kind, &values, names) - }) - .collect::, Snag>>()?; - batches.push(RecordBatch::try_new(schema.clone(), columns)?); - at += BATCH; + 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); } - // A result with no rows is still a result: it has a schema, and a - // reader that gets no batch at all cannot tell what the columns - // were. One empty batch says both. - if batches.is_empty() { - let columns = kinds + + 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(|kind| build(kind, &[], names)) - .collect::, Snag>>()?; - batches.push(RecordBatch::try_new(schema.clone(), columns)?); + .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)), + )) } - - let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); - Ok(FFI_ArrowArrayStream::new(Box::new(reader))) } -/// The type of a column, from the values in it. -fn infer<'a>(name: &str, values: impl Iterator) -> Result { - let mut kind = Kind::Null; - for (row, value) in values.enumerate() { - let found = kind_of(name, row, value)?; - let (held, arrived) = (kind.name(), found.name()); - kind = unify(kind, found).ok_or_else(|| { - Snag::Type(format!( - "column '{name}' mixes {held} and {arrived} at row {row}, and an Arrow column holds one type" - )) - })?; +impl RecordBatchReader for Slices { + fn schema(&self) -> SchemaRef { + self.schema.clone() } - Ok(kind) } -/// The type of one value, on its own. -fn kind_of(name: &str, row: usize, value: &Value) -> Result { - Ok(match value { - Value::Null => Kind::Null, - Value::Bool(_) => Kind::Bool, - Value::Int(_) => Kind::Int, - Value::Float(_) => Kind::Float, - Value::Str(_) => Kind::Str, - Value::Node { .. } => Kind::Node, - Value::Rel { .. } => Kind::Rel, - Value::Path(_) => Kind::Path, - Value::List(items) => { - let mut of = Kind::Null; - for item in items { - let found = kind_of(name, row, item)?; - let (held, arrived) = (of.name(), found.name()); - of = unify(of, found).ok_or_else(|| { - Snag::Type(format!( - "the list at row {row} of column '{name}' mixes {held} and {arrived}, and an Arrow list holds one type" - )) - })?; - } - Kind::List(Box::new(of)) - } - Value::Record(fields) => Kind::Record( - fields - .iter() - .map(|(field, value)| Ok((field.clone(), kind_of(name, row, value)?))) - .collect::, Snag>>()?, - ), - Value::Temporal(temporal) => match temporal { - Temporal::Date(_) => Kind::Date, - Temporal::LocalTime(_) => Kind::Time, - Temporal::LocalDatetime(_) => Kind::LocalDatetime, - Temporal::ZonedDatetime { offset, .. } => Kind::ZonedDatetime(*offset), - Temporal::Duration(DurationKind::YearMonth, _) => Kind::YearMonth, - Temporal::Duration(DurationKind::DayTime, _) => Kind::DayTime, - // 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. - Temporal::ZonedTime { .. } => { - return Err(Snag::Type(format!( - "row {row} of column '{name}' is 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. - Value::Graph(_) | Value::BindingTable(_) => { - return Err(Snag::Type(format!( - "row {row} of column '{name}' is a reference to a graph or a binding table, which Arrow has no type for" - ))); - } - // 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" - ))); - } - }) +/// 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))) } -/// The one type two types are both, or `None` when they are not. -fn unify(left: Kind, right: Kind) -> Option { - Some(match (left, right) { - (Kind::Null, other) | (other, Kind::Null) => other, - // The one widening: a projection that returns an integer for - // one row and a float for another means a number, and every - // reader of the column reads it as one. - (Kind::Int, Kind::Float) | (Kind::Float, Kind::Int) => Kind::Float, - // The first zoned value in the column names the zone. Later - // rows may have been written elsewhere, and they are the same - // instant either way, so this changes how a column prints and - // never what it holds. - (Kind::ZonedDatetime(offset), Kind::ZonedDatetime(_)) => Kind::ZonedDatetime(offset), - (Kind::List(left), Kind::List(right)) => Kind::List(Box::new(unify(*left, *right)?)), - (Kind::Record(left), Kind::Record(right)) => { - if left.len() != right.len() { - return None; - } - let mut fields = Vec::with_capacity(left.len()); - for ((name, left), (other, right)) in left.into_iter().zip(right) { - if name != other { - return None; - } - fields.push((name, unify(left, right)?)); +/// 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, + 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)), } - Kind::Record(fields) } - (left, right) if left == right => left, - _ => return None, + 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)?, }) } -/// One column's array, filled from the values in it. -fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result { - Ok(match kind { - Kind::Null => Arc::new(NullArray::new(values.len())), - Kind::Bool => Arc::new( +/// 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 { @@ -403,7 +420,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Int => Arc::new( + ColumnType::Int => Arc::new( values .iter() .map(|value| match value { @@ -412,7 +429,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Float => Arc::new( + ColumnType::Float => Arc::new( values .iter() .map(|value| match value { @@ -424,7 +441,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Str => Arc::new( + ColumnType::Str => Arc::new( values .iter() .map(|value| match value { @@ -433,7 +450,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Date => Arc::new( + ColumnType::Date => Arc::new( temporals(values) .map(|temporal| match temporal { Some(Temporal::Date(days)) => Some(*days), @@ -441,7 +458,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Time => Arc::new( + ColumnType::LocalTime => Arc::new( temporals(values) .map(|temporal| match temporal { Some(Temporal::LocalTime(nanos)) => Some(*nanos), @@ -449,7 +466,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::LocalDatetime => Arc::new( + ColumnType::LocalDatetime => Arc::new( temporals(values) .map(|temporal| match temporal { Some(Temporal::LocalDatetime(nanos)) => Some(*nanos), @@ -457,7 +474,7 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::ZonedDatetime(offset) => Arc::new( + ColumnType::ZonedDatetime { offset } => Arc::new( temporals(values) .map(|temporal| match temporal { Some(Temporal::ZonedDatetime { nanos, .. }) => Some(*nanos), @@ -466,29 +483,27 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result() .with_timezone(zone(*offset)), ), - Kind::YearMonth => { - let mut months = Vec::with_capacity(values.len()); - for (row, temporal) in temporals(values).enumerate() { - months.push(match temporal { + 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)) => { - // 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. - let count = i32::try_from(*count).map_err(|_| { - Snag::Value(format!( - "the duration at row {row} is {count} months, which is more than an Arrow interval holds" - )) - })?; - Some(IntervalMonthDayNano::new(count, 0, 0)) + counts.push(*count); + valid.push(true); } - _ => None, - }); + _ => { + counts.push(0); + valid.push(false); + } + } } - Arc::new(months.into_iter().collect::()) + Arc::new(IntervalMonthDayNanoArray::new( + ScalarBuffer::from(months(name, &counts)?), + Some(NullBuffer::from(valid)), + )) } - Kind::DayTime => Arc::new( + ColumnType::DayTime => Arc::new( temporals(values) .map(|temporal| match temporal { Some(Temporal::Duration(DurationKind::DayTime, nanos)) => Some(*nanos), @@ -496,10 +511,10 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result(), ), - Kind::Node => nodes(values, names)?, - Kind::Rel => rels(values, names)?, - Kind::Path => paths(values, names)?, - Kind::List(of) => { + 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()); @@ -514,15 +529,15 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result { + ColumnType::Record(fields) => { let mut children: Vec = Vec::with_capacity(fields.len()); - for (at, (_, kind)) in fields.iter().enumerate() { + for (at, (_, ty)) in fields.iter().enumerate() { let column: Vec<&Value> = values .iter() .map(|value| match value { @@ -530,17 +545,20 @@ fn build(kind: &Kind, values: &[&Value], names: &Names) -> Result &Value::Null, }) .collect(); - children.push(build(kind, &column, names)?); + children.push(build(name, ty, &column, names)?); } Arc::new(StructArray::try_new( - match kind.data_type() { + match data_type(name, ty)? { DataType::Struct(fields) => fields, - _ => unreachable!("a record is a struct"), + _ => return Err(mismatch(name, ty)), }, children, Some(present(values)), )?) } + ColumnType::ZonedTime { .. } | ColumnType::Graph | ColumnType::BindingTable => { + return Err(unsupported(name, ty)); + } }) } @@ -651,15 +669,25 @@ fn tables<'a>( /// 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(values: &[&Value], names: &Names) -> Result { +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 value in values { - if let Value::Path(elements) = value { - walked_nodes.extend(elements.iter().step_by(2)); - walked_rels.extend(elements.iter().skip(1).step_by(2)); + 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); diff --git a/src/conn.rs b/src/conn.rs index 9ebcb52..8e52472 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -549,6 +549,10 @@ 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. + /// + /// 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> { let py = slf.py(); needed(py, "pyarrow", "arrow")? diff --git a/tests/test_arrow.py b/tests/test_arrow.py index ac8b496..3399d3e 100644 --- a/tests/test_arrow.py +++ b/tests/test_arrow.py @@ -206,6 +206,37 @@ def test_a_result_with_no_rows_still_has_its_columns(loaded: zudb.Connection) -> assert table.schema.field("uid").type == pa.null() +@pytest.mark.parametrize( + "statement,arrow_type,answer", + [ + ("UNWIND [1, null, 3] AS v RETURN v", pa.int64(), [1, None, 3]), + ("UNWIND [1.5, null, 3.5] AS v RETURN v", pa.float64(), [1.5, None, 3.5]), + ("UNWIND [true, null, false] AS v RETURN v", pa.bool_(), [True, None, False]), + ("UNWIND ['a', null, 'ccc'] AS v RETURN v", pa.string(), ["a", None, "ccc"]), + ], +) +def test_a_gap_in_a_column_is_a_gap_in_the_bitmap( + empty: zudb.Connection, statement: str, arrow_type: pa.DataType, answer: list +) -> None: + # The values and the validity are two buffers, and a boolean column + # is the one where both of them are bits. A row that is null still + # occupies its cell in the values, which is what lets the column be + # handed over rather than rebuilt. + column = empty.execute(statement).to_arrow().column("v").combine_chunks() + assert column.type == arrow_type + assert column.null_count == 1 + assert column.to_pylist() == answer + + +def test_a_column_with_nothing_missing_carries_no_bitmap(empty: zudb.Connection) -> None: + # Absent rather than all ones, which is Arrow's own convention and + # the engine's: a reader of a column with nothing missing gets to + # skip the AND entirely. + column = empty.execute("UNWIND [1, 2, 3] AS v RETURN v").to_arrow().column("v") + assert column.null_count == 0 + assert column.combine_chunks().buffers()[0] is None + + def test_the_batches_are_the_same_rows(tmp_path: Path) -> None: rows = 70_000 zudb.load(tmp_path / "big.zu1", nodes="n", rels="r", columns={"uid": list(range(rows))}) diff --git a/tools/versus_duckdb.py b/tools/versus_duckdb.py new file mode 100644 index 0000000..6d34bfe --- /dev/null +++ b/tools/versus_duckdb.py @@ -0,0 +1,194 @@ +"""What a call costs here, against what the same call costs DuckDB. + +This is the script behind `docs/clients/duckdb.md` in the engine tree, +kept in the repository so the numbers in that page can be reproduced +rather than believed. A million rows of three columns, an integer, a +double and a short string, in a stored table on both sides, and the +fastest of five runs for each call, because what is being compared is +the code path and not the scheduler. + +One row of the output is not a comparison and is marked as such. +DuckDB's `execute` hands back a result nobody has read yet, so what it +takes is the cost of planning and of nothing else; ours has the whole +answer in memory by the time it returns. + +Run: python tools/versus_duckdb.py +Needs: pip install duckdb pyarrow pandas +""" + +from __future__ import annotations + +import shutil +import tempfile +import time +from pathlib import Path + +import duckdb +import zudb + +ROWS = 1_000_000 +RUNS = 5 + + +def best(run) -> float: + """The fastest of `RUNS`, in seconds.""" + fastest = float("inf") + for _ in range(RUNS): + at = time.perf_counter() + run() + fastest = min(fastest, time.perf_counter() - at) + return fastest + + +def read(call): + """A call that has to have read every row by the time it returns. + + The trap this closes is a handle: a call that hands back a lazy + reader has done no work, times at nought, and is not the same call + as one that has the answer. Anything that does not answer with + `ROWS` did not do what it was being timed for. + """ + + def run() -> None: + count = call() + if count != ROWS: + raise SystemExit(f"read {count} rows and not {ROWS}, so this is not the same call") + + return run + + +def rows() -> list[tuple[int, float, str]]: + return [(n, n * 1.5, f"s{n % 10_000:04d}") for n in range(ROWS)] + + +def build(directory: Path): + """Both databases, loaded with the same rows.""" + data = rows() + + zudb.load( + str(directory / "zu"), + nodes="row", + columns={ + "id": [row[0] for row in data], + "f": [row[1] for row in data], + "s": [row[2] for row in data], + }, + ) + zu = zudb.connect(str(directory / "zu")) + + duck = duckdb.connect(str(directory / "duck.db")) + duck.register("staged", _arrow(data)) + duck.execute("CREATE TABLE row AS SELECT * FROM staged") + duck.unregister("staged") + return zu, duck + + +def _arrow(data): + import pyarrow as pa + + return pa.table( + { + "id": pa.array([row[0] for row in data], pa.int64()), + "f": pa.array([row[1] for row in data], pa.float64()), + "s": pa.array([row[2] for row in data], pa.string()), + } + ) + + +def line(label: str, ours: float, theirs: float, note: str = "") -> None: + if note: + ratio = note + elif ours < theirs: + ratio = f"{theirs / ours:.1f}x faster" + else: + ratio = f"{ours / theirs:.1f}x slower" + print(f"| {label} | {ours * 1e3:.0f} ms | {theirs * 1e3:.0f} ms | {ratio} |") + + +def micro(label: str, ours: float, theirs: float) -> None: + """The same row, for a call small enough to read in microseconds.""" + if ours < theirs: + ratio = f"{theirs / ours:.1f}x faster" + else: + ratio = f"{ours / theirs:.1f}x slower" + print(f"| {label} | {ours * 1e6:.1f} us | {theirs * 1e6:.1f} us | {ratio} |") + + +def main() -> None: + directory = Path(tempfile.mkdtemp(prefix="zu-versus-")) + try: + zu, duck = build(directory) + query = "MATCH (r:row) RETURN r.id, r.f, r.s" + sql = "SELECT id, f, s FROM row" + + print("| call | zu | DuckDB | ratio |") + print("|---|---|---|---|") + line( + "execute, nothing read", + best(lambda: zu.execute(query)), + best(lambda: duck.execute(sql)), + note="not a comparison", + ) + line( + "execute and `fetchall`", + best(read(lambda: len(zu.execute(query).fetchall()))), + best(read(lambda: len(duck.execute(sql).fetchall()))), + ) + # DuckDB's `arrow()` hands back a `RecordBatchReader` that has + # read nothing, which times at nought and is not the same call. + # `to_arrow_table` is the one that has the table when it + # returns, which is what ours does. + line( + "execute and Arrow table", + best(read(lambda: zu.execute(query).to_arrow().num_rows)), + best(read(lambda: duck.execute(sql).to_arrow_table().num_rows)), + ) + line( + "execute and pandas", + best(read(lambda: len(zu.execute(query).to_pandas()))), + best(read(lambda: len(duck.execute(sql).df()))), + ) + + point = "MATCH (r:row) WHERE r.id = 500000 RETURN r.f" + point_sql = "SELECT f FROM row WHERE id = 500000" + + def swap(connection, table): + def run() -> None: + connection.register("frame", table) + connection.unregister("frame") + + return run + + table = _arrow(rows()) + numbers = table.select(["id", "f"]) + + print() + print("| call | zu | DuckDB | ratio |") + print("|---|---|---|---|") + micro( + "point read, whole call", + best(lambda: zu.execute(point).fetchall()), + best(lambda: duck.execute(point_sql).fetchall()), + ) + micro( + "`register` and `unregister`, numbers", + best(swap(zu, numbers)), + best(swap(duck, numbers)), + ) + # The same call with a string column in the frame, which is a + # different answer and is why both are printed: the bytes get + # a validation pass on the way in and the numbers do not. + micro( + "`register` and `unregister`, with strings", + best(swap(zu, table)), + best(swap(duck, table)), + ) + + zu.close() + duck.close() + finally: + shutil.rmtree(directory, ignore_errors=True) + + +if __name__ == "__main__": + main()