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()