From b981dd4420e364388322db15574247097a75fedc Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:02:55 +0700 Subject: [PATCH] an appender: rows in columns, a batch to a commit A statement writes one row at a time and commits each one, which is the wrong shape for loading data into a database that already exists. `load` is the right shape for one that does not exist yet and this is the right shape for one that does: rows go into per-column buffers in memory, and a flush turns the whole buffer into one commit. with conn.appender("person") as rows: for uid, name in enumerate(names): rows.append_row([uid, name]) The columns come from the table and not from the first row. An appender that read its shape off the row it was handed first would believe a wrong row and then refuse every right one after it, and its messages would name a column by number because that is all it would know. Reading the property directory at open costs one catalog read and buys a message that names the column, its type and the table. The engine's appender borrows the connection for as long as it lives, which is a promise a Python object cannot make, so this one buffers on this side and opens an engine appender for the length of a flush. That is a catalog read per flush against a commit and a fold that cost time proportional to the table, so it is not where a load spends its time. What it buys is an appender that can be held in a variable, passed to a function and closed by a `with` block. The GIL is released for the flush, and the buffer's own lock is taken after it is released rather than before: a lock waited for with the GIL held stops every other thread in the process for the length of the wait. That is the DX2 line about the GIL around an appender flush, and there is a test that counts the turns the main thread gets while 50,000 rows go in. Closing flushes, including on the way out of a block that raised. The Rust appender flushes when it is dropped for the same reason: a load that stopped partway is better served by its rows arriving than by them vanishing, and a caller who wants the other answer writes `discard()` and gets exactly it. An edge to a row that is not there is refused by the flush rather than left to the fold. The fold's refusal arrives once the write is durable, and the frame it leaves behind is refused again by every writer that opens the database afterwards, so one bad edge makes a database nobody can write to. The row counts are read at the flush and not at the open, so an edge to a row another appender wrote a moment ago is a good edge. The Python to column conversion moves to `buffer.rs`, since the loader and the appender both do it and only differ in what settles the type: the table when there is one, the first value when there is not. The loader keeps its one widening, a column of integers that meets a float, because nothing there has said what the column is. The appender has been told and refuses it. A load of a column of bytes is refused where it starts rather than written, because the store takes one and no statement can read one back yet. Numbers on this machine: 200,000 rows of an integer and a string take 11 ms to buffer through `append_rows` and 93 ms to flush, which is 1.9 million rows a second including the commit. Against `INSERT`, 2,000 rows take 32 seconds a row at a time and 28 ms through an appender. --- README.md | 18 +- python/zudb/__init__.py | 2 + python/zudb/_zudb.pyi | 41 +++ src/appender.rs | 611 ++++++++++++++++++++++++++++++++++++++++ src/buffer.rs | 356 +++++++++++++++++++++++ src/conn.rs | 23 +- src/error.rs | 26 +- src/lib.rs | 3 + src/load.rs | 234 +++------------ tests/test_appender.py | 461 ++++++++++++++++++++++++++++++ 10 files changed, 1569 insertions(+), 206 deletions(-) create mode 100644 src/appender.rs create mode 100644 src/buffer.rs create mode 100644 tests/test_appender.py diff --git a/README.md b/README.md index e6b271e..4778ad6 100644 --- a/README.md +++ b/README.md @@ -55,6 +55,22 @@ with zudb.connect("social.zu1", read_only=True) as conn: Edges name rows by position, counting from zero, because at load time a row has no other name. Columns may hold booleans, integers, floats, strings, dates, times, datetimes or durations, one kind to a column, and the GIL is released for the write. +## Adding to one that exists + +`load` writes a database and an appender grows one. Rows go into per-column buffers in memory and a flush turns the whole buffer into one commit, so a million rows cost one commit rather than a million. + +```python +with conn.appender("person") as rows: + for uid, name in enumerate(names): + rows.append_row([uid, name]) +``` + +A row is every column of the table, in the order the table declares them, and the columns come from the table rather than from the first row, so a value that does not belong in one is refused by the call that appended it and the message names the column. A rel table takes rows too, and a row of one is the two ends of an edge as offsets into the tables it runs between, which is how a Python program adds an edge at all. + +The `with` block closes, and closing flushes, including on the way out of a block that raised: a load that stopped partway is better served by its rows arriving than by them vanishing, and `discard()` is there for the caller who wants the other answer. A flush that fails keeps its rows, so what did not go in is still there to look at. + +On this machine 200,000 rows of an integer and a string take 11 ms to buffer through `append_rows` and 93 ms to flush, which is 1.9 million rows a second including the commit. The same rows one call at a time cost 32 ms of buffering, since a call is a call. Against `INSERT`, 2,000 rows take 32 seconds a row at a time and 28 ms through an appender, and the gap widens with the table because every `INSERT` is a commit and a fold. + ## Reading a result as columns A result is rows to iterate and columns to hand to something else. The columns go out over the Arrow C Data Interface, so pyarrow, pandas and polars each read the same buffers and none of them gets a Python object per cell. @@ -77,7 +93,7 @@ The stub is checked against the module it describes in CI: griffe reads the stub ## What works today -The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, stubs inside the wheel with a gate that keeps them true, and the GIL released around every statement, every load and every copy out. `register` and the interrupt are next, and each one lands with the tests that say it works. +The list above is what this client is for. What it does so far is the core of it: `connect`, `execute` and `sql` with named parameters, results that iterate and fetch, values as Python objects both ways including dates, times, datetimes and durations, `Node`, `Rel` and `Path` as classes, `load` for building a graph with edges in it, an appender for growing one, every condition as an exception class carrying its code, its position and its documentation link, results as Arrow columns and as pandas and polars frames, stubs inside the wheel with a gate that keeps them true, and the GIL released around every statement, every load and every copy out. `register` and the interrupt are next, and each one lands with the tests that say it works. ## Wheels diff --git a/python/zudb/__init__.py b/python/zudb/__init__.py index a1f652c..6b0d606 100644 --- a/python/zudb/__init__.py +++ b/python/zudb/__init__.py @@ -15,6 +15,7 @@ from __future__ import annotations from ._zudb import ( + Appender, Connection, Duration, Node, @@ -43,6 +44,7 @@ "connect", "load", "Connection", + "Appender", "Result", "Node", "Rel", diff --git a/python/zudb/_zudb.pyi b/python/zudb/_zudb.pyi index 503aec8..166f08e 100644 --- a/python/zudb/_zudb.pyi +++ b/python/zudb/_zudb.pyi @@ -69,6 +69,9 @@ class Connection: def sql(self, statement: str, params: Mapping[str, Value] | None = None) -> Result: """The same call, named for the way it reads in a notebook.""" + def appender(self, table: str) -> Appender: + """Opens an appender on `table`, for loading rows into a database that already exists.""" + def close(self) -> None: """Closes the connection and frees what it held.""" @@ -76,6 +79,44 @@ class Connection: def __exit__(self, *_exception: object) -> bool: ... def __repr__(self) -> str: ... +class Appender: + """Rows on their way into a table, buffered until they are flushed.""" + + @property + def table(self) -> str: + """The table this appender writes to.""" + + @property + def buffered(self) -> int: + """Rows buffered and not yet written.""" + + @property + def committed(self) -> int: + """Rows this appender has committed, across every flush.""" + + @property + def closed(self) -> bool: + """Whether this appender has been closed.""" + + def append_row(self, row: Iterable[Value]) -> None: + """Appends one row, which is a sequence of one value per column of the table.""" + + def append_rows(self, rows: Iterable[Iterable[Value]]) -> int: + """Appends every row of an iterable of rows.""" + + def flush(self) -> int: + """Writes every buffered row and makes it readable.""" + + def discard(self) -> int: + """Throws away what is buffered and answers how many rows that was.""" + + def close(self) -> int: + """Flushes what is left and answers how many rows this appender committed in all.""" + + def __enter__(self) -> Appender: ... + def __exit__(self, *_exception: object) -> bool: ... + def __repr__(self) -> str: ... + class Result: """The rows a statement gave back.""" diff --git a/src/appender.rs b/src/appender.rs new file mode 100644 index 0000000..ce385d2 --- /dev/null +++ b/src/appender.rs @@ -0,0 +1,611 @@ +//! Appending rows to a table that already exists. +//! +//! `INSERT` is the wrong shape for loading data. Every row is parsed, +//! bound, planned and committed, and the commit is the expensive part, +//! so a million rows is a million commits and the load is dominated by +//! durability work nobody asked for. `load` is the right shape for a +//! database that does not exist yet, and this is the right shape for +//! one that does: rows go into per-column buffers in memory, and a +//! flush turns the whole buffer into one commit. +//! +//! ```python +//! with conn.appender("person") as rows: +//! for uid, name in enumerate(names): +//! rows.append_row([uid, name]) +//! ``` +//! +//! A row is every column of the table, in the order the table declares +//! them, and a column is a position rather than a name: naming the +//! columns per row would cost a lookup per value on the one path where +//! per-value cost is the whole story, and a loader knows its own column +//! order. +//! +//! The engine's appender borrows the connection for as long as it +//! lives, which is a promise a Python object cannot make, so this one +//! buffers here and opens an engine appender for the length of a flush. +//! That is a catalog read per flush, against a commit and a fold that +//! cost time proportional to the table, so it is not where a load +//! spends its time. What it buys is an appender that can be held in a +//! variable, passed to a function and closed by a `with` block, which +//! is what a Python caller expects of it. + +use std::sync::{Mutex, MutexGuard}; + +use pyo3::prelude::*; +use pyo3::types::PyTuple; +use zudb::zu1::catalog::Catalog; +use zudb::{Field, ZuError}; + +use crate::buffer::{Column, Mismatch, type_name}; +use crate::conn::Connection; +use crate::error::{programming, to_py_err}; + +/// Rows on their way into a table, buffered until they are flushed. +/// +/// Take one with `Connection.appender`, append rows to it, and close +/// it. What is buffered is columnar and typed from the table's own +/// columns, read when the appender opened, so a value that does not +/// belong in a column is refused by the call that appended it rather +/// than at the flush that would have carried it, and the message names +/// the column it did not fit. +#[pyclass(module = "zudb")] +pub struct Appender { + /// The connection this writes through, held rather than borrowed: + /// an appender whose connection was collected would be a buffer + /// with nowhere to go. + conn: Py, + #[pyo3(get)] + table: String, + state: Mutex, +} + +/// One column of the table, and what has been buffered for it. +struct Buffer { + name: String, + values: Column, + /// The node table whose rows this column names, for the two columns + /// of a rel table and for nothing else: a row of one is an offset + /// into the table the edge runs from and an offset into the table + /// it runs to. A negative offset is no row of anything and is + /// refused where it was appended; whether the row is there at all + /// is a question only the flush can answer, since the table may be + /// being appended to at the same time. + ends: Option, +} + +/// What is buffered, and how much of it has gone in. +struct State { + /// One buffer per column of the table, in the order the table + /// declares them, built when the appender opened. A flush empties + /// these and keeps them, since the next batch is the same shape as + /// the last. + cols: Vec, + /// Rows buffered and not yet written, which is every column's + /// length and is kept beside them so that a table with no columns + /// can still answer for itself. + buffered: u64, + /// Rows this appender has committed, across every flush. + committed: u64, + open: bool, +} + +/// What can go wrong with the GIL down, where there is no way to build +/// a Python exception yet. +enum Snag { + Closed, + Locked, + Finished, + Engine(ZuError), + /// A row the engine's appender would not take, named by where it + /// sits in the batch, because that is the row the caller can go and + /// look at. + Row(u64, ZuError), + /// An edge to a row that is not there, caught before the write + /// rather than after it. + Offset { + row: u64, + offset: i64, + table: String, + rows: u64, + }, +} + +impl Snag { + fn raise(self, py: Python<'_>) -> PyErr { + match self { + Snag::Closed => programming( + py, + "the connection this appender writes through is closed, so its \ + buffered rows have nowhere to go", + ), + Snag::Locked => programming( + py, + "this appender was left locked by a panic, and what it holds \ + cannot be trusted to be a rectangle any more", + ), + Snag::Finished => programming( + py, + "this appender is closed, and a closed appender has already \ + written everything it was given", + ), + Snag::Engine(err) => to_py_err(py, err), + // The engine reports the value and the column; which row of + // the batch it was is the part only this side knows, and it + // is the part that says where to look. + Snag::Row(at, err) => { + let raised = to_py_err(py, err); + programming(py, &format!("row {at} of this batch: {}", raised.value(py))) + } + Snag::Offset { + row, + offset, + table, + rows, + } => pyo3::exceptions::PyValueError::new_err(format!( + "row {row} of this batch joins row {offset} of '{table}', which has {rows} \ + rows in it, so the rows an edge joins have to be written before the edge is" + )), + } + } +} + +#[pymethods] +impl Appender { + /// Appends one row, which is a sequence of one value per column of + /// the table, in the order the table declares them. + /// + /// The values go into memory and nothing else happens, so this is a + /// conversion and a push per column. A row of the wrong width, or + /// with a value that does not fit the column, is refused with + /// nothing of it kept, so the appender is still usable once the + /// caller has fixed the row. + fn append_row(&self, py: Python<'_>, row: &Bound<'_, PyAny>) -> PyResult<()> { + let mut state = self.writable(py)?; + state.append(&self.table, row) + } + + /// Appends every row of an iterable of rows. + /// + /// The same thing in a loop, and worth a call of its own because it + /// is one lock and one attribute lookup for the batch rather than + /// one per row. A row that is refused stops the call where it was + /// refused and the rows before it stay buffered: nothing here is a + /// transaction until the flush, and throwing away work the caller + /// can keep would not make it one. + fn append_rows(&self, py: Python<'_>, rows: &Bound<'_, PyAny>) -> PyResult { + let mut state = self.writable(py)?; + let mut taken = 0; + for row in rows.try_iter()? { + state.append(&self.table, &row?)?; + taken += 1; + } + Ok(taken) + } + + /// Writes every buffered row and makes it readable, and answers how + /// many rows this appender has committed in all. + /// + /// One commit, whatever the buffer holds, with the GIL released for + /// it: the values are sealed into the file, one frame naming them + /// is synced to the log, and the fold that follows puts them where + /// every query looks. On return the buffer is empty and the rows + /// are there. A flush with nothing buffered touches no file, so a + /// loader can flush on a timer without writing empty commits. + /// + /// A flush that fails keeps its rows, so that what did not go in is + /// still there to be looked at and tried again. + fn flush(&self, py: Python<'_>) -> PyResult { + let held = self.conn.bind(py).borrow(); + let conn: &Connection = &held; + // Everything from here down is Rust over buffers that are + // already typed, so the GIL is not needed for any of it, and a + // flush is where a load spends its time. The buffer's own lock + // is taken down here too, and not above: a lock waited for with + // the GIL held is a lock that stops every other thread in the + // process for as long as the wait. + py.detach(|| -> Result { + let mut state = self.state.lock().map_err(|_| Snag::Locked)?; + if !state.open { + return Err(Snag::Finished); + } + state.write(conn, &self.table) + }) + .map_err(|snag| snag.raise(py)) + } + + /// Rows buffered and not yet written. + #[getter] + fn buffered(&self, py: Python<'_>) -> PyResult { + Ok(self.locked(py)?.buffered) + } + + /// Rows this appender has committed, across every flush. + #[getter] + fn committed(&self, py: Python<'_>) -> PyResult { + Ok(self.locked(py)?.committed) + } + + /// Throws away what is buffered and answers how many rows that was. + /// + /// The way out of a load that went wrong halfway. A caller who has + /// noticed that the rows are wrong wants them gone, and closing + /// would write them. Rows an earlier flush committed are committed, + /// and this does not reach them. + fn discard(&self, py: Python<'_>) -> PyResult { + let mut state = self.writable(py)?; + let dropped = state.buffered; + state.empty(); + Ok(dropped) + } + + /// Flushes what is left and answers how many rows this appender + /// committed in all. + /// + /// Closing twice is not an error and writes nothing the second + /// time, because a `with` block that closed early would otherwise + /// fail on the way out. + fn close(&self, py: Python<'_>) -> PyResult { + let held = self.conn.bind(py).borrow(); + let conn: &Connection = &held; + py.detach(|| -> Result { + let mut state = self.state.lock().map_err(|_| Snag::Locked)?; + if !state.open { + return Ok(state.committed); + } + let committed = state.write(conn, &self.table)?; + state.open = false; + Ok(committed) + }) + .map_err(|snag| snag.raise(py)) + } + + /// Whether this appender has been closed. + #[getter] + fn closed(&self, py: Python<'_>) -> PyResult { + Ok(!self.locked(py)?.open) + } + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + /// Closes on the way out, which flushes, and does it whether the + /// block ended well or badly. + /// + /// A block that raised is the case worth thinking about, and the + /// rows still go in. That is the answer the Rust appender gives by + /// flushing when it is dropped, for the same reason: a loader that + /// stopped partway is better served by its rows arriving than by + /// them vanishing, and a caller who wants the other answer writes + /// `discard()` and gets exactly it. An error the flush raises + /// carries the block's own exception as its context, so neither of + /// the two is lost. + #[pyo3(signature = (*_exception))] + fn __exit__(&self, py: Python<'_>, _exception: &Bound<'_, PyTuple>) -> PyResult { + self.close(py)?; + // False, so an exception raised inside the block carries on out + // of it. + Ok(false) + } + + fn __repr__(&self, py: Python<'_>) -> PyResult { + let state = self.locked(py)?; + let closed = if state.open { "" } else { ", closed" }; + Ok(format!( + "", + self.table, state.buffered, state.committed + )) + } +} + +impl Appender { + /// Opens an appender on `table`, which is a node table or a rel + /// table of the graph this connection reads. + /// + /// An engine appender is opened here and dropped, purely to find + /// out whether it can be opened at all: a table nothing declares, a + /// column that holds a null, a table a keyed rel table is built + /// over, and a connection that is read-only are all refused here + /// rather than at the first flush. A caller about to buffer a + /// million rows wants to hear about them now. + pub fn open(py: Python<'_>, conn: Py, table: &str) -> PyResult { + let cols = { + let held = conn.bind(py).borrow(); + let held: &Connection = &held; + py.detach(|| -> Result, Snag> { + let mut engine = held.inner.lock().map_err(|_| Snag::Closed)?; + let engine = engine.as_mut().ok_or(Snag::Closed)?; + engine.appender(table).map(drop).map_err(Snag::Engine)?; + shape(engine, table) + }) + .map_err(|snag| snag.raise(py))? + }; + Ok(Appender { + conn, + table: table.to_string(), + state: Mutex::new(State { + cols, + buffered: 0, + committed: 0, + open: true, + }), + }) + } + + /// The buffers, for a call that only reads them. + /// + /// The lock is taken with the GIL held, which a flush never does. + /// Waiting for it here waits for at most one conversion of one row, + /// and a flush that is under way has already let go of it by the + /// time it needs the GIL back. + fn locked(&self, py: Python<'_>) -> PyResult> { + self.state.lock().map_err(|_| Snag::Locked.raise(py)) + } + + /// The buffers, for a call that adds to them, which a closed + /// appender has no business doing. + fn writable(&self, py: Python<'_>) -> PyResult> { + let state = self.locked(py)?; + if !state.open { + return Err(Snag::Finished.raise(py)); + } + Ok(state) + } +} + +impl State { + /// One row into the buffers, or nothing at all. + fn append(&mut self, table: &str, row: &Bound<'_, PyAny>) -> PyResult<()> { + let width = self.cols.len(); + let mut at = 0; + for value in row.try_iter()? { + let value = match value { + Ok(value) => value, + Err(err) => return self.refuse(at, err), + }; + if at == width { + return self.refuse( + at, + pyo3::exceptions::PyValueError::new_err(format!( + "this row carries more than the {width} values '{table}' takes: {}", + self.names() + )), + ); + } + if let Err(err) = self.cols[at].take(&value, table, at) { + return self.refuse(at, err); + } + at += 1; + } + if at != width { + return self.refuse( + at, + pyo3::exceptions::PyValueError::new_err(format!( + "this row carries {at} values and '{table}' takes {width}: {}", + self.names() + )), + ); + } + self.buffered += 1; + Ok(()) + } + + /// The columns of the table, named, for a message about a row that + /// is the wrong shape. A caller who miscounted wants to see what + /// the count was supposed to be made of. + fn names(&self) -> String { + self.cols + .iter() + .map(|column| column.name.as_str()) + .collect::>() + .join(", ") + } + + /// Takes back the values a refused row managed to write, so that a + /// refused row is a row that never happened rather than half of one + /// nobody can find. A ragged buffer would be refused by the ingest + /// at the flush, a long way from the row that caused it. + fn refuse(&mut self, written: usize, err: PyErr) -> PyResult<()> { + for column in self.cols.iter_mut().take(written) { + column.values.pop(); + } + Err(err) + } + + /// The write itself, with the GIL already down. + fn write(&mut self, conn: &Connection, table: &str) -> Result { + if self.buffered == 0 { + return Ok(self.committed); + } + let rows = self.buffered; + { + let mut engine = conn.inner.lock().map_err(|_| Snag::Closed)?; + let engine = engine.as_mut().ok_or(Snag::Closed)?; + self.reachable(engine, rows)?; + let mut appender = engine.appender(table).map_err(Snag::Engine)?; + // One vector, refilled per row rather than allocated per + // row, which over a million rows is one allocation rather + // than a million. The fields borrow the buffers, which is + // what keeps a string column to one copy on the way in and + // one on the way out rather than three. + let mut row: Vec> = Vec::with_capacity(self.cols.len()); + for at in 0..rows as usize { + row.clear(); + row.extend(self.cols.iter().map(|column| column.values.field(at))); + appender + .append_row(&row[..]) + .map_err(|err| Snag::Row(at as u64, err))?; + } + appender.close().map_err(Snag::Engine)?; + } + self.empty(); + self.committed += rows; + Ok(self.committed) + } + + /// Every edge joins two rows that are there, checked against the + /// row counts as they stand at the flush. + /// + /// This is the flush's own check and not the engine's, because the + /// engine's comes too late: an edge to a row that is not there is + /// refused when the write is folded into the graph, which is after + /// the write is durable, and the frame it leaves behind is refused + /// again by every writer that opens the database afterwards. Caught + /// here, the batch is refused and the file is untouched. + /// + /// The counts are read at the flush and not when the appender + /// opened, because the rows a later edge names may be written by an + /// earlier flush of another appender on the same connection, and an + /// edge to a row that arrived in the meantime is a good edge. + fn reachable(&self, engine: &mut zudb::Connection, rows: u64) -> Result<(), Snag> { + if self.cols.iter().all(|column| column.ends.is_none()) { + return Ok(()); + } + let catalog = catalog(engine)?; + for column in &self.cols { + let Some(end) = column.ends else { continue }; + let Some(node) = catalog.node_by_id(end) else { + continue; + }; + for at in 0..rows as usize { + let Field::Int(offset) = column.values.field(at) else { + continue; + }; + if offset as u64 >= node.node_count { + return Err(Snag::Offset { + row: at as u64, + offset, + table: node.name.clone(), + rows: node.node_count, + }); + } + } + } + Ok(()) + } + + fn empty(&mut self) { + self.cols + .iter_mut() + .for_each(|column| column.values.clear()); + self.buffered = 0; + } +} + +impl Buffer { + /// One value into this column, or the reason it does not go there. + /// + /// The column's own name is in the message, and its position too, + /// because a row is written by position and read by name and a + /// caller who has them the wrong way round needs both to see it. + fn take(&mut self, value: &Bound<'_, PyAny>, table: &str, at: usize) -> PyResult<()> { + if self.ends.is_some() + && let Ok(offset) = value.extract::() + && offset < 0 + { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "value {at} of this row is {offset}, and column '{}' of '{table}' holds row \ + offsets, which count from zero", + self.name + ))); + } + self.values.push(value).map_err(|why| match why { + Mismatch::Wanted(holds) => pyo3::exceptions::PyTypeError::new_err(format!( + "value {at} of this row is of type '{}' and column '{}' of '{table}' holds {holds}", + type_name(value), + self.name + )), + Mismatch::Python(err) => err, + }) + } +} + +/// The catalog as the file has it, which is not always the one the +/// session is holding. +/// +/// A session reloads its catalog when a statement runs, and an appender +/// is not a statement: a flush commits its rows and folds them in +/// without the session hearing about it, so the row counts the session +/// remembers are the counts from the last statement. Loading it costs a +/// read of a block chain, once per appender opened and once per flush of +/// a rel table, which is nothing beside the commit either of them is +/// about to do. +fn catalog(engine: &mut zudb::Connection) -> Result { + Catalog::load(engine.session_mut().file_mut()).map_err(Snag::Engine) +} + +/// The columns of the table an appender was opened on, in the order it +/// declares them. +/// +/// A node table's columns are the ones the property store holds, with +/// the types it holds them as, which is what the engine's appender +/// checks a row against. A rel table has no property columns: a row of +/// one is the two ends of an edge, as offsets into the tables it runs +/// between, so those are the two columns and they are named for what +/// they are. +/// +/// Read here rather than left to the flush so that a value that does +/// not belong in a column is refused by the call that appended it. A +/// row is refused a million rows before the flush that would have +/// carried it, and the message names the column rather than guessing at +/// it from the values that came before. +fn shape(engine: &mut zudb::Connection, table: &str) -> Result, Snag> { + let catalog = catalog(engine)?; + if let Some(rel) = catalog.rel_by_name(table) { + let ends = [rel.from, rel.to]; + let named = |id: u32, fallback: &str| { + catalog + .node_tables() + .iter() + .find(|node| node.id == id) + .map_or_else(|| fallback.to_string(), |node| node.name.clone()) + }; + // Named for the tables the edge runs between, since that is + // what a row of a rel table is and there is nothing else to + // call the two columns. + return Ok(vec![ + Buffer { + name: format!("from {}", named(ends[0], "the source table")), + values: Column::Int(Vec::new()), + ends: Some(ends[0]), + }, + Buffer { + name: format!("to {}", named(ends[1], "the destination table")), + values: Column::Int(Vec::new()), + ends: Some(ends[1]), + }, + ]); + } + let id = catalog + .node_by_name(table) + .map(|node| node.id) + .ok_or_else(|| { + Snag::Engine(ZuError::InvalidArgument(format!( + "no node table or rel table '{table}'" + ))) + })?; + let directory = zudb::zu1::props::load_props(engine.session_mut().file_mut(), id) + .map_err(Snag::Engine)? + .ok_or_else(|| { + Snag::Engine(ZuError::InvalidArgument(format!( + "'{table}' stores no properties, so it has no columns to append to" + ))) + })?; + directory + .columns + .iter() + .map(|column| { + Ok(Buffer { + name: column.name.clone(), + values: Column::for_type(&column.ty).ok_or_else(|| { + Snag::Engine(ZuError::InvalidArgument(format!( + "column '{}' of '{table}' holds {}, which this engine cannot yet \ + append to", + column.name, column.ty + ))) + })?, + ends: None, + }) + }) + .collect() +} diff --git a/src/buffer.rs b/src/buffer.rs new file mode 100644 index 0000000..af0479d --- /dev/null +++ b/src/buffer.rs @@ -0,0 +1,356 @@ +//! Python values, buffered as the columns the engine stores. +//! +//! Two callers want the same thing here and would otherwise want it +//! twice. A load reads whole columns and writes them once; an appender +//! reads whole rows and writes them a batch at a time. Both of them end +//! up holding a vector per column in the shape the property store +//! keeps it in, and both of them have to decide what a Python object +//! is on the way in. +//! +//! What a column holds is settled by the table it is being written to +//! when there is one, which is the appender's case, and by its first +//! value when there is not, which is the loader's: a load writes a +//! database that does not exist yet, so nothing but the values can say +//! what the columns are. There is no null either way. A column that +//! holds one cannot be loaded or appended to, so a null here could only +//! ever be refused, and refusing it where it is named is better than +//! refusing it at the end of a million rows. + +use pyo3::prelude::*; +use pyo3::types::{PyBool, PyBytes, PyDate, PyDateTime, PyDelta, PyTime}; +use zu_common::temporal::days_from_civil; +use zu_common::{DurationKind, FloatBits, IntBits, LogicalType, Temporal}; +use zudb::Field; + +use crate::value::{Duration, clock_nanos}; + +/// One column's values, in the shape the property store wants them. +/// +/// Owned rather than borrowed from the caller's lists, because a +/// Python list holds objects and the store holds numbers: there is +/// nothing here to borrow. The arms are the storage arms, so a flush +/// or a load hands a buffer straight over with no pass to convert it. +pub enum Column { + Int(Vec), + Float(Vec), + Bool(Vec), + /// Kept as strings rather than as bytes, because the store wants + /// the bytes and the appender wants the `&str`, and a `String` + /// lends out either without a copy. + Str(Vec), + Bytes(Vec>), + Date(Vec), + LocalTime(Vec), + LocalDatetime(Vec), + Duration(DurationKind, Vec), +} + +/// Why a value did not go in. +/// +/// A column that wanted something else says what it holds and lets the +/// caller word the rest, because a load knows the column's name and a +/// row of an appender knows its position, and neither message reads +/// well in the other's place. Anything else is Python's own error and +/// is passed along as it is. +pub enum Mismatch { + Wanted(&'static str), + Python(PyErr), +} + +impl From for Mismatch { + fn from(err: PyErr) -> Mismatch { + Mismatch::Python(err) + } +} + +impl Column { + /// The buffer a column of this declared type appends into, or + /// `None` for a type the ingest path cannot carry. + /// + /// The match is on the exact declared type and not on its family, + /// because that is what the ingest checks: it compares the stored + /// column's type against the type its values claim, so an `INT32` + /// column or a `VARCHAR(20)` one has no buffer here even though its + /// bits would fit the same lane. This is the engine appender's own + /// table, kept in step with it, because a buffer it would refuse is + /// better refused before a million rows go into it. + pub fn for_type(ty: &LogicalType) -> Option { + Some(match ty { + LogicalType::Int { + signed: true, + bits: IntBits::B64, + precision: None, + } => Column::Int(Vec::new()), + LogicalType::Bool => Column::Bool(Vec::new()), + LogicalType::Float { + bits: FloatBits::B64, + precision: None, + } => Column::Float(Vec::new()), + LogicalType::Date => Column::Date(Vec::new()), + LogicalType::LocalTime => Column::LocalTime(Vec::new()), + LogicalType::LocalDatetime => Column::LocalDatetime(Vec::new()), + LogicalType::Duration(kind) => Column::Duration(*kind, Vec::new()), + LogicalType::Str { + min: None, + max: None, + fixed: false, + } => Column::Str(Vec::new()), + LogicalType::Bytes { + min: None, + max: None, + fixed: false, + } => Column::Bytes(Vec::new()), + _ => return None, + }) + } + + /// The column a first value starts, with that value already in it, + /// or `None` for a value no column here holds. + /// + /// The order matters twice. Bool before int, since in Python every + /// bool is an int and a column of `True` is not a column of ones. + /// Datetime before date, since a datetime is a date and reading one + /// as the other would throw the time away. + pub fn start(value: &Bound<'_, PyAny>) -> PyResult> { + if let Ok(b) = value.cast::() { + return Ok(Some(Column::Bool(vec![b.is_true()]))); + } + if let Ok(s) = value.extract::() { + return Ok(Some(Column::Str(vec![s]))); + } + if let Ok(b) = value.cast::() { + return Ok(Some(Column::Bytes(vec![b.as_bytes().to_vec()]))); + } + if let Ok(n) = value.extract::() { + return Ok(Some(Column::Int(vec![n as u64]))); + } + if let Ok(f) = value.extract::() { + return Ok(Some(Column::Float(vec![f]))); + } + if let Ok(d) = value.extract::() { + return Ok(Some(if d.months == 0 { + Column::Duration(DurationKind::DayTime, vec![d.nanoseconds]) + } else { + Column::Duration(DurationKind::YearMonth, vec![d.months]) + })); + } + if let Ok(delta) = value.cast::() { + return Ok(Some(Column::Duration( + DurationKind::DayTime, + vec![delta_nanos(delta)?], + ))); + } + if let Ok(dt) = value.cast::() { + return Ok(Some(Column::LocalDatetime(vec![datetime_nanos(dt)?]))); + } + if let Ok(date) = value.cast::() { + return Ok(Some(Column::Date(vec![date_days(date)?]))); + } + if let Ok(time) = value.cast::() { + return Ok(Some(Column::LocalTime(vec![clock_nanos(time.as_any())?]))); + } + Ok(None) + } + + /// Takes one more value, or says what this column holds instead. + /// + /// The order of the arms is the order of [`Column::start`] and for + /// the same two reasons: a bool is an int in Python and a datetime + /// is a date, so each of those pairs has to be told apart before it + /// is read. + pub fn push(&mut self, value: &Bound<'_, PyAny>) -> Result<(), Mismatch> { + let holds = self.holds(); + let wanted = || Mismatch::Wanted(holds); + match self { + Column::Bool(v) => v.push(value.cast::().map_err(|_| wanted())?.is_true()), + Column::Int(v) => { + if value.cast::().is_ok() { + return Err(wanted()); + } + v.push(value.extract::().map_err(|_| wanted())? as u64); + } + Column::Float(v) => v.push(value.extract::().map_err(|_| wanted())?), + Column::Str(v) => v.push(value.extract::().map_err(|_| wanted())?), + Column::Bytes(v) => v.push( + value + .cast::() + .map_err(|_| wanted())? + .as_bytes() + .to_vec(), + ), + Column::LocalDatetime(v) => { + let dt = value.cast::().map_err(|_| wanted())?; + v.push(datetime_nanos(dt)?); + } + Column::Date(v) => { + if value.cast::().is_ok() { + return Err(wanted()); + } + let date = value.cast::().map_err(|_| wanted())?; + v.push(date_days(date)?); + } + Column::LocalTime(v) => { + let time = value.cast::().map_err(|_| wanted())?; + v.push(clock_nanos(time.as_any())?); + } + Column::Duration(kind, v) => { + let kind = *kind; + v.push(duration_count(kind, value).ok_or_else(wanted)??); + } + } + Ok(()) + } + + /// Takes one more value, and turns a column of integers into a + /// column of floats rather than refusing a float. + /// + /// The one mixture worth widening, and only where the column's type + /// is not already settled by a table: a list written by hand as + /// `[1, 2.5]` is a column of floats and refusing it would be + /// pedantry. The values already in it are exactly what they were, + /// up to the point where an integer stops fitting in a float, which + /// is a different problem and a rarer one. + pub fn widening_push(&mut self, value: &Bound<'_, PyAny>) -> Result<(), Mismatch> { + if let Column::Int(v) = self + && value.cast::().is_err() + && value.extract::().is_err() + && let Ok(f) = value.extract::() + { + let mut widened: Vec = v.iter().map(|&n| n as i64 as f64).collect(); + widened.push(f); + *self = Column::Float(widened); + return Ok(()); + } + self.push(value) + } + + /// What this column holds, for the message when it was handed + /// something else. Plural, because it is the column that is being + /// described and not the value. + pub fn holds(&self) -> &'static str { + match self { + Column::Int(_) => "integers", + Column::Float(_) => "floats", + Column::Bool(_) => "booleans", + Column::Str(_) => "strings", + Column::Bytes(_) => "byte strings", + Column::Date(_) => "dates", + Column::LocalTime(_) => "times", + Column::LocalDatetime(_) => "datetimes", + Column::Duration(DurationKind::YearMonth, _) => "year-month durations", + Column::Duration(DurationKind::DayTime, _) => "day-time durations", + } + } + + /// How many values are in it, which is how many rows the table has + /// if this is the first column and a refusal if it is not. + pub fn len(&self) -> usize { + match self { + Column::Int(v) => v.len(), + Column::Float(v) => v.len(), + Column::Bool(v) => v.len(), + Column::Str(v) => v.len(), + Column::Bytes(v) => v.len(), + Column::Date(v) => v.len(), + Column::LocalTime(v) | Column::LocalDatetime(v) => v.len(), + Column::Duration(_, v) => v.len(), + } + } + + /// Drops the value written last, which is how a refused row takes + /// back the fields it managed to write before the one that failed. + pub fn pop(&mut self) { + match self { + Column::Int(v) => drop(v.pop()), + Column::Float(v) => drop(v.pop()), + Column::Bool(v) => drop(v.pop()), + Column::Str(v) => drop(v.pop()), + Column::Bytes(v) => drop(v.pop()), + Column::Date(v) => drop(v.pop()), + Column::LocalTime(v) | Column::LocalDatetime(v) => drop(v.pop()), + Column::Duration(_, v) => drop(v.pop()), + } + } + + pub fn clear(&mut self) { + match self { + Column::Int(v) => v.clear(), + Column::Float(v) => v.clear(), + Column::Bool(v) => v.clear(), + Column::Str(v) => v.clear(), + Column::Bytes(v) => v.clear(), + Column::Date(v) => v.clear(), + Column::LocalTime(v) | Column::LocalDatetime(v) => v.clear(), + Column::Duration(_, v) => v.clear(), + } + } + + /// One value of this column as the engine's appender takes it. + /// + /// A field borrows rather than owning, which is the point of it on + /// a string column: the buffer already holds the bytes and the + /// appender is about to copy them into its own, so lending them is + /// the difference between one copy per row and two. + pub fn field(&self, row: usize) -> Field<'_> { + match self { + Column::Int(v) => Field::Int(v[row] as i64), + Column::Float(v) => Field::Float(v[row]), + Column::Bool(v) => Field::Bool(v[row]), + Column::Str(v) => Field::Str(&v[row]), + Column::Bytes(v) => Field::Bytes(&v[row]), + Column::Date(v) => Field::Temporal(Temporal::Date(v[row])), + Column::LocalTime(v) => Field::Temporal(Temporal::LocalTime(v[row])), + Column::LocalDatetime(v) => Field::Temporal(Temporal::LocalDatetime(v[row])), + Column::Duration(kind, v) => Field::Temporal(Temporal::Duration(*kind, v[row])), + } + } +} + +/// A duration for a column that is already one of the two kinds, or +/// `None` for a duration of the other kind, which is the one place the +/// two do not mix: a column of months has no room for a count of +/// nanoseconds and the other way about. +fn duration_count(kind: DurationKind, value: &Bound<'_, PyAny>) -> Option> { + if let Ok(d) = value.extract::() { + return match kind { + DurationKind::YearMonth if d.months != 0 || d.nanoseconds == 0 => Some(Ok(d.months)), + DurationKind::DayTime if d.months == 0 => Some(Ok(d.nanoseconds)), + _ => None, + }; + } + match (kind, value.cast::()) { + (DurationKind::DayTime, Ok(delta)) => Some(delta_nanos(delta)), + _ => None, + } +} + +/// The name a Python type goes by, for a message about a value that was +/// the wrong one. +pub fn type_name(value: &Bound<'_, PyAny>) -> String { + value + .get_type() + .getattr("__name__") + .and_then(|name| name.extract::()) + .unwrap_or_else(|_| "unknown".to_string()) +} + +pub fn date_days(date: &Bound<'_, PyDate>) -> PyResult { + Ok(days_from_civil( + date.getattr("year")?.extract()?, + date.getattr("month")?.extract()?, + date.getattr("day")?.extract()?, + )) +} + +pub fn datetime_nanos(dt: &Bound<'_, PyDateTime>) -> PyResult { + const NANOS_PER_DAY: i64 = 86_400 * 1_000_000_000; + let days = date_days(dt.as_any().cast::()?)?; + Ok(i64::from(days) * NANOS_PER_DAY + clock_nanos(dt.as_any())?) +} + +pub fn delta_nanos(delta: &Bound<'_, PyDelta>) -> PyResult { + let days: i64 = delta.getattr("days")?.extract()?; + let seconds: i64 = delta.getattr("seconds")?.extract()?; + let micros: i64 = delta.getattr("microseconds")?.extract()?; + Ok(days * 86_400 * 1_000_000_000 + seconds * 1_000_000_000 + micros * 1_000) +} diff --git a/src/conn.rs b/src/conn.rs index 6100992..92668bc 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -16,6 +16,7 @@ use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use zudb::query::{QueryResult, Value}; use zudb::{Config, Database}; +use crate::appender::Appender; use crate::columns; use crate::error::{closed, to_py_err}; use crate::value::{Names, from_py, to_py}; @@ -35,7 +36,15 @@ const STREAM: &CStr = c"arrow_array_stream"; pub struct Connection { /// `None` once closed, which is what makes a second `close()` do /// nothing and a statement after one an error rather than a crash. - inner: Mutex>, + /// + /// Reachable from the appender, which writes through the same + /// connection and takes this same lock. One rule holds for every + /// caller of it: take it with the GIL already released. A thread + /// that waited for it holding the GIL would stop every other + /// thread in the process for the length of somebody else's + /// statement, and would deadlock against the thread inside that + /// statement, which needs the GIL back to return. + pub(crate) inner: Mutex>, #[pyo3(get)] path: PathBuf, #[pyo3(get)] @@ -98,6 +107,18 @@ impl Connection { self.execute(py, statement, params) } + /// Opens an appender on `table`, for loading rows into a database + /// that already exists. + /// + /// A statement writes one row at a time and commits each one, which + /// is the wrong shape for a million rows. An appender buffers them + /// in columns and writes each batch as one commit. The table has to + /// be there already, since an appender adds rows to a table and + /// does not make one. + fn appender(slf: Py, py: Python<'_>, table: &str) -> PyResult { + Appender::open(py, slf, table) + } + /// Closes the connection and frees what it held. /// /// Doing it twice is not an error, because a `with` block that diff --git a/src/error.rs b/src/error.rs index 38f4135..6268b55 100644 --- a/src/error.rs +++ b/src/error.rs @@ -120,20 +120,28 @@ fn severity(severity: zudb::Severity) -> &'static str { } } -/// What a call raises when it is made on a connection that is closed. +/// A mistake the program made, as the class for one. /// -/// A `ValueError` would be wrong and a segfault would be worse: the -/// caller made a mistake, in Python, and `zudb.ProgrammingError` is -/// the class for a mistake a program made rather than a condition the -/// engine raised. -pub fn closed(py: Python<'_>, what: &str) -> PyErr { +/// A `ValueError` would be wrong and a segfault would be worse: +/// `zudb.ProgrammingError` is the class for something the caller did +/// rather than for a condition the engine raised, and it is what a +/// call made on an object that is closed, or handed something it +/// cannot do anything with, raises. +pub fn programming(py: Python<'_>, message: &str) -> PyErr { match errors(py).and_then(|errors| { let class: Bound<'_, PyType> = errors.getattr("ProgrammingError")?.cast_into()?; - Ok(PyErr::from_value(class.call1((format!( - "{what} is closed, so there is nothing left to run a statement on" - ),))?)) + Ok(PyErr::from_value(class.call1((message,))?)) }) { Ok(raised) => raised, Err(broken) => broken, } } + +/// What a statement raises when it is run on a connection that is +/// closed. +pub fn closed(py: Python<'_>, what: &str) -> PyErr { + programming( + py, + &format!("{what} is closed, so there is nothing left to run a statement on"), + ) +} diff --git a/src/lib.rs b/src/lib.rs index 7ce89c2..862dd03 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,6 +12,8 @@ //! that work. What it owes in return is the ABI's semantics, and the //! conformance corpus is what says whether it paid. +mod appender; +mod buffer; mod columns; mod conn; mod error; @@ -51,6 +53,7 @@ fn _zudb(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(load::load, module)?)?; module.add_class::()?; module.add_class::()?; + module.add_class::()?; module.add_class::()?; module.add_class::()?; module.add_class::()?; diff --git a/src/load.rs b/src/load.rs index cc7ef0f..8badc89 100644 --- a/src/load.rs +++ b/src/load.rs @@ -17,49 +17,13 @@ use std::path::PathBuf; use pyo3::prelude::*; -use pyo3::types::{PyBool, PyDate, PyDateTime, PyDelta, PyDict, PyTime}; -use zu_common::DurationKind; -use zu_common::temporal::days_from_civil; +use pyo3::types::PyDict; use zudb::zu1::file::Zu1File; use zudb::zu1::graph::bulk_load_keyed; use zudb::zu1::props::{PropValues, store_props}; +use crate::buffer::{Column, Mismatch, type_name}; use crate::error::to_py_err; -use crate::value::{Duration, clock_nanos}; - -/// One column, reduced to the vector the property store keeps it in. -/// -/// Owned rather than borrowed from the caller's lists, because a -/// Python list holds objects and the store holds numbers: there is -/// nothing here to borrow. Which arm a column is comes from its first -/// value, and every value after it has to be that arm too. -enum Column { - Int(Vec), - Float(Vec), - Bool(Vec), - Str(Vec>), - Date(Vec), - LocalTime(Vec), - LocalDatetime(Vec), - Duration(DurationKind, Vec), -} - -impl Column { - /// How many values are in it, which is how many rows the table has - /// if this is the first column and a refusal if it is not. - fn len(&self) -> usize { - match self { - Column::Int(v) => v.len(), - Column::Float(v) => v.len(), - Column::Bool(v) => v.len(), - Column::Str(v) => v.len(), - Column::Date(v) => v.len(), - Column::LocalTime(v) => v.len(), - Column::LocalDatetime(v) => v.len(), - Column::Duration(_, v) => v.len(), - } - } -} /// Writes a new database at `path` and answers what went into it. /// @@ -119,22 +83,24 @@ pub fn load( pairs.dedup(); bulk_load_keyed(&mut db, nodes, rels, rows, &pairs, None)?; if !built.is_empty() { - // The store wants a slice of slices for a string column, - // which a `Vec>` is not, so the row borrows are - // built first and handed over after. - let strings: Vec> = built + // The store wants a slice of slices for a column of strings + // or of bytes, which a vector of either is not, so the row + // borrows are built first and handed over after. + let runs: Vec> = built .iter() .map(|(_, column)| match column { - Column::Str(v) => v.iter().map(Vec::as_slice).collect(), + Column::Str(v) => v.iter().map(String::as_bytes).collect(), + Column::Bytes(v) => v.iter().map(Vec::as_slice).collect(), _ => Vec::new(), }) .collect(); let props: Vec<(&str, PropValues<'_>)> = built .iter() - .zip(&strings) - .map(|((name, column), strings)| { + .zip(&runs) + .map(|((name, column), runs)| { let values = match column { - Column::Str(_) => PropValues::Str(strings), + Column::Str(_) => PropValues::Str(runs), + Column::Bytes(_) => PropValues::Bytes(runs), Column::Int(v) => PropValues::Int(v), Column::Float(v) => PropValues::Float(v), Column::Bool(v) => PropValues::Bool(v), @@ -174,6 +140,16 @@ fn build(columns: Option<&Bound<'_, PyDict>>) -> PyResult> )); } let column = column(&name, &values)?; + // The store takes a column of bytes and every statement that + // reads one back refuses it, so a load that wrote one would be + // writing data the caller cannot get at again. Refused here + // until the read side catches up, at which point this goes and + // nothing else has to change. + if matches!(column, Column::Bytes(_)) { + return Err(pyo3::exceptions::PyTypeError::new_err(format!( + "column '{name}' holds byte strings, and no statement can read one back yet, so a load will not write a column of them" + ))); + } if let Some((first, had)) = built.first().map(|(name, column)| (name, column.len())) && column.len() != had { @@ -190,69 +166,29 @@ fn build(columns: Option<&Bound<'_, PyDict>>) -> PyResult> /// One column, read out of a sequence of Python objects. /// /// The first value settles what the column is and every value after it -/// has to be the same thing. There is no null: a column that holds one -/// cannot be loaded this way, so a null here could only be refused, -/// and refusing it where it is named is better than refusing it at the -/// end of a million rows. +/// has to agree, which is [`Column`]'s rule and is the appender's rule +/// too. What this adds is the column's name, because a message about a +/// column is worth leading with the name of it. fn column(name: &str, values: &Bound<'_, PyAny>) -> PyResult { let mut column: Option = None; for (row, value) in values.try_iter()?.enumerate() { let value = value?; - let mismatch = |want: &str| { - let got = value - .get_type() - .getattr("__name__") - .and_then(|name| name.extract::()) - .unwrap_or_else(|_| "unknown".to_string()); - pyo3::exceptions::PyTypeError::new_err(format!( - "column '{name}' holds {want} and row {row} is of type '{got}'" - )) - }; match column.as_mut() { - None => column = Some(started(name, row, &value)?), - // Bool before int, since in Python every bool is an int - // and a column of `True` is not a column of ones. - Some(Column::Bool(v)) => v.push( - value - .cast::() - .map_err(|_| mismatch("booleans"))? - .is_true(), - ), - Some(Column::Int(v)) => { - if value.cast::().is_ok() { - return Err(mismatch("integers")); - } - v.push(value.extract::().map_err(|_| mismatch("integers"))? as u64); - } - Some(Column::Float(v)) => { - v.push(value.extract::().map_err(|_| mismatch("floats"))?) - } - Some(Column::Str(v)) => v.push( - value - .extract::() - .map_err(|_| mismatch("strings"))? - .into_bytes(), - ), - // Datetime before date, since a datetime is a date and - // reading one as the other would throw the time away. - Some(Column::LocalDatetime(v)) => { - let dt = value - .cast::() - .map_err(|_| mismatch("datetimes"))?; - v.push(datetime_nanos(dt)?); - } - Some(Column::Date(v)) => { - if value.cast::().is_ok() { - return Err(mismatch("dates")); - } - let date = value.cast::().map_err(|_| mismatch("dates"))?; - v.push(date_days(date)?); - } - Some(Column::LocalTime(v)) => { - let time = value.cast::().map_err(|_| mismatch("times"))?; - v.push(clock_nanos(time.as_any())?); + Some(column) => column.widening_push(&value).map_err(|why| match why { + Mismatch::Wanted(holds) => pyo3::exceptions::PyTypeError::new_err(format!( + "column '{name}' holds {holds} and row {row} is of type '{}'", + type_name(&value) + )), + Mismatch::Python(err) => err, + })?, + None => { + column = Some(Column::start(&value)?.ok_or_else(|| { + pyo3::exceptions::PyTypeError::new_err(format!( + "column '{name}' starts at row {row} with a value of type '{}', and a loaded column holds booleans, integers, floats, strings, dates, times, datetimes or durations", + type_name(&value) + )) + })?); } - Some(Column::Duration(kind, v)) => v.push(duration_count(*kind, &value, mismatch)?), } } column.ok_or_else(|| { @@ -262,98 +198,6 @@ fn column(name: &str, values: &Bound<'_, PyAny>) -> PyResult { }) } -/// The column a first value starts, with that value already in it. -fn started(name: &str, row: usize, value: &Bound<'_, PyAny>) -> PyResult { - if let Ok(b) = value.cast::() { - return Ok(Column::Bool(vec![b.is_true()])); - } - if let Ok(s) = value.extract::() { - return Ok(Column::Str(vec![s.into_bytes()])); - } - if let Ok(n) = value.extract::() { - return Ok(Column::Int(vec![n as u64])); - } - if let Ok(f) = value.extract::() { - return Ok(Column::Float(vec![f])); - } - if let Ok(d) = value.extract::() { - return Ok(if d.months == 0 { - Column::Duration(DurationKind::DayTime, vec![d.nanoseconds]) - } else { - Column::Duration(DurationKind::YearMonth, vec![d.months]) - }); - } - if let Ok(delta) = value.cast::() { - return Ok(Column::Duration( - DurationKind::DayTime, - vec![delta_nanos(delta)?], - )); - } - if let Ok(dt) = value.cast::() { - return Ok(Column::LocalDatetime(vec![datetime_nanos(dt)?])); - } - if let Ok(date) = value.cast::() { - return Ok(Column::Date(vec![date_days(date)?])); - } - if let Ok(time) = value.cast::() { - return Ok(Column::LocalTime(vec![clock_nanos(time.as_any())?])); - } - let got = value - .get_type() - .getattr("__name__") - .and_then(|name| name.extract::()) - .unwrap_or_else(|_| "unknown".to_string()); - Err(pyo3::exceptions::PyTypeError::new_err(format!( - "column '{name}' starts at row {row} with a value of type '{got}', and a loaded column holds booleans, integers, floats, strings, dates, times, datetimes or durations" - ))) -} - -/// A duration for a column that is already one of the two kinds, which -/// is the one place the two do not mix: a column of months has no room -/// for a count of nanoseconds and the other way about. -fn duration_count( - kind: DurationKind, - value: &Bound<'_, PyAny>, - mismatch: impl Fn(&str) -> PyErr, -) -> PyResult { - let want = match kind { - DurationKind::YearMonth => "year-month durations", - DurationKind::DayTime => "day-time durations", - }; - if let Ok(d) = value.extract::() { - return match kind { - DurationKind::YearMonth if d.months != 0 || d.nanoseconds == 0 => Ok(d.months), - DurationKind::DayTime if d.months == 0 => Ok(d.nanoseconds), - _ => Err(mismatch(want)), - }; - } - match (kind, value.cast::()) { - (DurationKind::DayTime, Ok(delta)) => delta_nanos(delta), - _ => Err(mismatch(want)), - } -} - -fn date_days(date: &Bound<'_, PyDate>) -> PyResult { - Ok(days_from_civil( - date.getattr("year")?.extract()?, - date.getattr("month")?.extract()?, - date.getattr("day")?.extract()?, - )) -} - -fn datetime_nanos(dt: &Bound<'_, PyDateTime>) -> PyResult { - const NANOS_PER_DAY: i64 = 86_400 * 1_000_000_000; - let days = date_days(dt.as_any().cast::()?)?; - Ok(i64::from(days) * NANOS_PER_DAY + clock_nanos(dt.as_any())?) -} - -fn delta_nanos(delta: &Bound<'_, PyDelta>) -> PyResult { - let days: i64 = delta.getattr("days")?.extract()?; - let seconds: i64 = delta.getattr("seconds")?.extract()?; - let micros: i64 = delta.getattr("microseconds")?.extract()?; - Ok(days * 86_400 * 1_000_000_000 + seconds * 1_000_000_000 + micros * 1_000) -} - /// The edge list, as the pairs of row numbers it is. /// /// An edge naming a row the table has not got is refused here rather diff --git a/tests/test_appender.py b/tests/test_appender.py new file mode 100644 index 0000000..2cf252a --- /dev/null +++ b/tests/test_appender.py @@ -0,0 +1,461 @@ +"""Rows on their way into a table that already exists. + +An appender is the fast way in and the only way that is not a statement, +so these check the three things that makes it: that the rows arrive and +read back as themselves, that a row which cannot mean anything is +refused where it was appended rather than at the flush, and that a +refusal leaves the appender and the database exactly as they were. + +The last of those is the one worth the most tests. A load runs for a +long time, and a caller who is a million rows in and has just been told +that row 999,999 is wrong wants the other 999,998 and a database that +still opens. +""" + +from __future__ import annotations + +import datetime +import gc +import threading +import time +from pathlib import Path + +import pytest +import zudb + + +@pytest.fixture +def graph(tmp_path: Path) -> zudb.Connection: + """Three people and the edges between two of them, open for writing. + + Loaded rather than inserted because a rel table is made by a load + and by nothing else, and half of what an appender is for is adding + edges to one. + """ + path = tmp_path / "graph.zu1" + zudb.load( + path, + nodes="person", + rels="knows", + columns={"uid": [10, 20, 30], "name": ["ada", "grace", "kay"]}, + edges=[(0, 1)], + ) + conn = zudb.connect(path) + yield conn + conn.close() + + +def names(conn: zudb.Connection) -> list[str]: + return [name for (name,) in conn.execute("MATCH (p:person) RETURN p.name AS name")] + + +def edges(conn: zudb.Connection) -> list[tuple[str, str]]: + return list( + conn.execute("MATCH (a:person)-[:knows]->(b:person) RETURN a.name AS a, b.name AS b") + ) + + +def test_a_row_appended_and_flushed_is_a_row_you_can_query(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.append_row([40, "hopper"]) + assert app.flush() == 1 + assert names(graph) == ["ada", "grace", "kay", "hopper"] + + +def test_nothing_is_written_until_the_flush(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.append_row([40, "hopper"]) + assert app.buffered == 1 + assert app.committed == 0 + assert names(graph) == ["ada", "grace", "kay"] + app.close() + + +def test_a_flush_empties_the_buffer_and_leaves_the_appender_open(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.append_row([40, "hopper"]) + app.flush() + assert (app.buffered, app.committed, app.closed) == (0, 1, False) + app.append_row([50, "liskov"]) + assert app.flush() == 2 + assert names(graph) == ["ada", "grace", "kay", "hopper", "liskov"] + + +def test_append_rows_takes_a_batch_and_says_how_many_it_took(graph: zudb.Connection) -> None: + app = graph.appender("person") + assert app.append_rows([[40, "hopper"], [50, "liskov"]]) == 2 + assert app.buffered == 2 + assert app.close() == 2 + assert names(graph) == ["ada", "grace", "kay", "hopper", "liskov"] + + +def test_append_rows_reads_an_iterator_as_happily_as_a_list(graph: zudb.Connection) -> None: + rows = ([uid, f"p{uid}"] for uid in (40, 50, 60)) + app = graph.appender("person") + assert app.append_rows(rows) == 3 + assert app.close() == 3 + + +def test_a_with_block_flushes_on_the_way_out(graph: zudb.Connection) -> None: + with graph.appender("person") as app: + app.append_row([40, "hopper"]) + assert app.closed + assert names(graph) == ["ada", "grace", "kay", "hopper"] + + +def test_a_block_that_raised_still_writes_the_rows_it_managed(graph: zudb.Connection) -> None: + # The Rust appender flushes when it is dropped and this is the same + # answer for the same reason: a load that stopped partway is better + # served by its rows arriving than by them vanishing, and a caller + # who wants the other answer has `discard()`. + with pytest.raises(RuntimeError, match="halfway"): + with graph.appender("person") as app: + app.append_row([40, "hopper"]) + raise RuntimeError("stopped halfway") + assert names(graph) == ["ada", "grace", "kay", "hopper"] + + +def test_discard_throws_away_what_is_buffered(graph: zudb.Connection) -> None: + with graph.appender("person") as app: + app.append_rows([[40, "hopper"], [50, "liskov"]]) + assert app.discard() == 2 + assert app.buffered == 0 + assert names(graph) == ["ada", "grace", "kay"] + + +def test_discard_does_not_reach_what_a_flush_committed(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.append_row([40, "hopper"]) + app.flush() + app.append_row([50, "liskov"]) + assert app.discard() == 1 + assert app.close() == 1 + assert names(graph) == ["ada", "grace", "kay", "hopper"] + + +def test_closing_twice_writes_nothing_the_second_time(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.append_row([40, "hopper"]) + assert app.close() == 1 + assert app.close() == 1 + assert app.closed + + +def test_a_closed_appender_takes_no_more_rows(graph: zudb.Connection) -> None: + app = graph.appender("person") + app.close() + with pytest.raises(zudb.ProgrammingError, match="closed appender"): + app.append_row([40, "hopper"]) + with pytest.raises(zudb.ProgrammingError, match="closed appender"): + app.flush() + with pytest.raises(zudb.ProgrammingError, match="closed appender"): + app.discard() + + +def test_a_flush_with_nothing_buffered_commits_nothing(graph: zudb.Connection) -> None: + app = graph.appender("person") + assert app.flush() == 0 + assert app.flush() == 0 + assert names(graph) == ["ada", "grace", "kay"] + app.close() + + +def test_a_row_of_the_wrong_width_is_refused_by_name(graph: zudb.Connection) -> None: + app = graph.appender("person") + with pytest.raises(ValueError, match="carries 1 values and 'person' takes 2: uid, name"): + app.append_row([40]) + with pytest.raises(ValueError, match="more than the 2 values 'person' takes"): + app.append_row([40, "hopper", "extra"]) + app.close() + + +def test_a_refused_row_leaves_the_buffer_a_rectangle(graph: zudb.Connection) -> None: + # The half of the row that went in has to come back out, or the + # columns are different lengths and the flush is the one that finds + # out, a long way from the row that did it. + app = graph.appender("person") + app.append_row([40, "hopper"]) + with pytest.raises(TypeError): + app.append_row([50, 50]) + assert app.buffered == 1 + app.append_row([50, "liskov"]) + assert app.close() == 2 + assert names(graph) == ["ada", "grace", "kay", "hopper", "liskov"] + + +def test_a_value_the_column_does_not_hold_is_refused_where_it_was_appended( + graph: zudb.Connection, +) -> None: + app = graph.appender("person") + with pytest.raises( + TypeError, match="value 0 of this row is of type 'str' and column 'uid' of 'person'" + ): + app.append_row(["forty", "hopper"]) + app.close() + + +def test_the_table_says_what_a_column_holds_and_not_the_first_row(graph: zudb.Connection) -> None: + # A first row that was wrong used to settle the shape and then + # refuse every right row after it. The columns come from the table, + # so the wrong row is the one refused. + app = graph.appender("person") + with pytest.raises(TypeError, match="column 'uid' of 'person' holds integers"): + app.append_row(["hopper", 40]) + app.append_row([40, "hopper"]) + assert app.close() == 1 + + +def test_an_integer_column_refuses_a_float(graph: zudb.Connection) -> None: + # A load widens a column of integers when it meets a float, because + # nothing there has said what the column is. Here the table has + # said, and 4.5 is not an integer. + app = graph.appender("person") + with pytest.raises(TypeError, match="'float' and column 'uid' of 'person' holds integers"): + app.append_row([4.5, "hopper"]) + app.close() + + +def test_a_bool_is_not_an_integer(graph: zudb.Connection) -> None: + app = graph.appender("person") + with pytest.raises(TypeError, match="'bool' and column 'uid' of 'person' holds integers"): + app.append_row([True, "hopper"]) + app.close() + + +def test_every_type_a_column_holds_goes_in_and_comes_back(tmp_path: Path) -> None: + path = tmp_path / "types.zu1" + zudb.load( + path, + nodes="thing", + columns={ + "n": [1], + "f": [1.5], + "b": [True], + "s": ["one"], + "d": [datetime.date(2020, 1, 1)], + "t": [datetime.time(1, 2, 3)], + "ts": [datetime.datetime(2020, 1, 1, 1, 2, 3)], + "dur": [datetime.timedelta(days=1)], + "ym": [zudb.Duration(months=3)], + }, + ) + row = [ + 2, + 2.5, + False, + "two", + datetime.date(2022, 3, 4), + datetime.time(5, 6, 7), + datetime.datetime(2022, 3, 4, 5, 6, 7), + datetime.timedelta(minutes=30), + zudb.Duration(months=5), + ] + with zudb.connect(path) as conn: + with conn.appender("thing") as app: + app.append_row(row) + got = conn.execute( + "MATCH (t:thing) WHERE t.n = 2 RETURN t.n, t.f, t.b, t.s, t.d, t.t, t.ts, t.dur, t.ym" + ).fetchone() + # A duration comes back as a `Duration`, since Python's own type + # cannot hold the year-month half of one, so the `timedelta` that + # went in is the one value that does not read back as itself. + assert list(got) == [*row[:7], zudb.Duration(nanoseconds=30 * 60 * 1_000_000_000), row[8]] + + +def test_a_float_column_takes_an_integer(tmp_path: Path) -> None: + # The one widening the other way round, and the only one: every + # integer a Python program is likely to hand a float column is a + # float exactly, and refusing 7 for a column of scores would be + # pedantry rather than safety. + path = tmp_path / "scores.zu1" + zudb.load(path, nodes="person", columns={"score": [36.5]}) + with zudb.connect(path) as conn: + with conn.appender("person") as app: + app.append_row([7]) + assert conn.execute("MATCH (p:person) RETURN p.score AS s").fetchall() == [(36.5,), (7.0,)] + + +def test_an_appender_on_a_rel_table_joins_the_graph(graph: zudb.Connection) -> None: + with graph.appender("knows") as rels: + rels.append_row([1, 2]) + assert edges(graph) == [("ada", "grace"), ("grace", "kay")] + + +def test_a_rel_row_is_two_offsets_and_a_negative_one_is_no_row(graph: zudb.Connection) -> None: + with graph.appender("knows") as rels: + with pytest.raises(ValueError, match="row offsets, which count from zero"): + rels.append_row([0, -1]) + assert rels.buffered == 0 + + +def test_an_edge_to_a_row_that_is_not_there_is_refused_before_it_is_written( + graph: zudb.Connection, +) -> None: + # Refused by the flush and not by the fold that comes after it. The + # fold's refusal arrives once the write is durable, and the frame it + # leaves behind is refused again by every writer that opens the + # database afterwards, which is a database nobody can write to over + # one bad edge. + rels = graph.appender("knows") + rels.append_row([0, 99]) + with pytest.raises(ValueError, match="joins row 99 of 'person', which has 3 rows"): + rels.flush() + assert rels.buffered == 1 + rels.discard() + rels.close() + assert edges(graph) == [("ada", "grace")] + + +def test_the_database_still_opens_after_an_edge_that_was_refused(tmp_path: Path) -> None: + path = tmp_path / "graph.zu1" + zudb.load(path, nodes="person", rels="knows", columns={"uid": [1, 2]}, edges=[(0, 1)]) + with zudb.connect(path) as conn: + rels = conn.appender("knows") + rels.append_row([0, 99]) + with pytest.raises(ValueError): + rels.close() + rels.discard() + rels.close() + with zudb.connect(path) as conn: + assert conn.execute("MATCH ()-[r:knows]->() RETURN count(r) AS n").fetchone() == (1,) + + +def test_an_edge_can_name_a_row_a_flush_wrote_a_moment_ago(graph: zudb.Connection) -> None: + # The row counts are read at the flush and not when the appender + # opened, so a rel appender held across a load of nodes writes the + # edges to them rather than refusing every one. + rels = graph.appender("knows") + with graph.appender("person") as people: + people.append_row([40, "hopper"]) + rels.append_row([0, 3]) + rels.close() + assert edges(graph) == [("ada", "grace"), ("ada", "hopper")] + + +def test_a_table_nothing_declares_has_no_appender(graph: zudb.Connection) -> None: + with pytest.raises(zudb.ProgrammingError, match="no node table or rel table 'cities'"): + graph.appender("cities") + + +def test_a_read_only_connection_has_no_appender(tmp_path: Path) -> None: + path = tmp_path / "graph.zu1" + zudb.load(path, nodes="person", columns={"uid": [1]}) + with zudb.connect(path, read_only=True) as conn: + with pytest.raises(zudb.ProgrammingError, match="the connection is read-only"): + conn.appender("person") + + +def test_a_closed_connection_has_nowhere_to_put_the_rows(tmp_path: Path) -> None: + path = tmp_path / "graph.zu1" + zudb.load(path, nodes="person", columns={"uid": [1]}) + conn = zudb.connect(path) + app = conn.appender("person") + app.append_row([2]) + conn.close() + with pytest.raises(zudb.ProgrammingError, match="connection this appender writes through"): + app.flush() + + +def test_an_appender_keeps_its_connection_alive(tmp_path: Path) -> None: + # The connection is held and not borrowed, so an appender handed + # back by a function that opened one is an appender that still + # works. A borrowed one would be a buffer with nowhere to go. + path = tmp_path / "graph.zu1" + zudb.load(path, nodes="person", columns={"uid": [1]}) + + def opened() -> zudb.Appender: + return zudb.connect(path).appender("person") + + app = opened() + gc.collect() + app.append_row([2]) + assert app.close() == 1 + with zudb.connect(path, read_only=True) as conn: + assert conn.execute("MATCH (p:person) RETURN count(p) AS n").fetchone() == (2,) + + +def test_the_repr_says_what_it_is_holding(graph: zudb.Connection) -> None: + app = graph.appender("person") + assert app.table == "person" + assert repr(app) == "" + app.append_row([40, "hopper"]) + assert repr(app) == "" + app.flush() + assert repr(app) == "" + app.close() + assert repr(app) == "" + + +def test_two_threads_appending_to_one_appender_lose_nothing(graph: zudb.Connection) -> None: + app = graph.appender("person") + + def run(start: int) -> None: + for uid in range(start, start + 200): + app.append_row([uid, f"p{uid}"]) + + threads = [threading.Thread(target=run, args=(base,)) for base in (1000, 2000, 3000)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=120) + assert not thread.is_alive(), "a thread is still waiting for the appender" + assert app.buffered == 600 + assert app.close() == 600 + assert len(names(graph)) == 603 + + +def test_python_keeps_running_while_a_flush_does(graph: zudb.Connection) -> None: + ticks = 0 + done = threading.Event() + app = graph.appender("person") + app.append_rows([[uid, f"p{uid}"] for uid in range(50_000)]) + + def run() -> None: + app.close() + done.set() + + worker = threading.Thread(target=run) + worker.start() + while not done.is_set(): + ticks += 1 + worker.join(timeout=120) + assert not worker.is_alive() + # A GIL held for the length of the flush would leave this loop no + # turns at all rather than thousands of them. + assert ticks > 1000, f"the main thread only got {ticks} turns" + + +#: Rows for the comparison against `INSERT`, and few enough that the +#: `INSERT` half finishes in a few seconds. It is the slow half by three +#: orders of magnitude, and it gets slower as the table grows, because +#: every row of it is a commit and a fold. +COMPARED = 200 + + +def test_appending_beats_inserting_by_the_margin_that_makes_it_worth_having( + tmp_path: Path, +) -> None: + rows = [(uid, f"p{uid}") for uid in range(1, COMPARED)] + + with zudb.connect(tmp_path / "inserted.zu1") as conn: + conn.execute("INSERT (p:person {uid: 0, name: 'seed'})") + started = time.perf_counter() + for uid, name in rows: + conn.execute("INSERT (p:person {uid: $u, name: $n})", {"u": uid, "n": name}) + inserting = time.perf_counter() - started + + with zudb.connect(tmp_path / "appended.zu1") as conn: + conn.execute("INSERT (p:person {uid: 0, name: 'seed'})") + started = time.perf_counter() + with conn.appender("person") as app: + app.append_rows(rows) + appending = time.perf_counter() - started + assert conn.execute("MATCH (p:person) RETURN count(p) AS n").fetchone() == (COMPARED,) + + # Measured at about 150 times on this machine at this row count and + # rising with it, since one commit is one commit however many rows + # it carries. The gate is 20, which is the number that says the + # appender is still batching rather than the number it hits. + assert inserting > 20 * appending, ( + f"{COMPARED} rows: {inserting * 1000:.0f} ms inserted, {appending * 1000:.0f} ms appended" + )