From b4ba6b977c4fe3272f3fce0972e22b3efd720820 Mon Sep 17 00:00:00 2001 From: Tam Nguyen Duc <1218621+tamnd@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:17:50 +0700 Subject: [PATCH] a result as columns, and the frames built on them A result is rows to iterate and columns to hand to something else. The rows were there already. This is the columns: `__arrow_c_stream__` on `Result`, and `to_arrow`, `to_pandas`, `to_polars` and `record_batches` written on top of it. The interface is the PyCapsule one, so the capsule is the whole of it and the four methods are conveniences. `pyarrow.table(result)` and `polars.DataFrame(result)` read a result with nothing here doing anything, which is the point of a protocol every reader implements. `to_pandas` asks for Arrow-backed dtypes, which is what pandas 3 wants anyway and what keeps a string column from becoming a column of Python strings on the way in. `QueryResult` carries no column types, only values, so the type of a column is inferred from what is in it. One type to a column, because that is what Arrow holds. Integers beside floats widen to floats, which is the only mixture that is not refused: anything else says which two types it mixed and at which row. Nulls do not decide anything and a column of nothing but nulls is Arrow's null type, which is also what an empty result's columns are. Nodes, rels and paths go across as structs, and a path is a struct of two lists. A year-month duration goes as a month-day-nano interval with the days and nanoseconds zero, because pyarrow has no array class for the year-month one and raises a KeyError on the type id. A time with an offset is refused, because Arrow has no type for one. Batches are 65,536 rows. The copy runs with the GIL released. On this machine 300,000 rows across three columns take 44 ms as Arrow against 67 ms as Python objects, and a single integer column takes 13.8 ms against 44.5 ms, which is where the difference actually lives: a Python object per cell is the cost, and not the query. Table names are borrowed rather than cloned per row, which took the node column from 27 ms to 17.5 ms over 300,000 rows. 33 tests, skipped when pyarrow is not installed, and the pandas and polars ones skip on their own. The wheel still depends on nothing. --- .github/workflows/ci.yml | 2 +- Cargo.lock | 434 ++++++++++++++++++++++++- Cargo.toml | 5 + README.md | 20 +- pyproject.toml | 16 +- src/columns.rs | 673 +++++++++++++++++++++++++++++++++++++++ src/conn.rs | 97 +++++- src/lib.rs | 1 + src/value.rs | 26 +- tests/test_arrow.py | 276 ++++++++++++++++ 10 files changed, 1535 insertions(+), 15 deletions(-) create mode 100644 src/columns.rs create mode 100644 tests/test_arrow.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index debdc2c..87bfde2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -45,6 +45,6 @@ jobs: # the suite imports is what a person who runs `pip install zudb` # gets: the extension out of a wheel, the package out of # site-packages, and nothing resolved out of the checkout. - - run: pip install . + - run: pip install ".[all]" - run: pip install pytest - run: pytest diff --git a/Cargo.lock b/Cargo.lock index 1aa6476..cd2fc5e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,20 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.5" @@ -20,6 +34,169 @@ dependencies = [ "libc", ] +[[package]] +name = "arrow" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.17.1", + "libc", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64", + "chrono", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-data" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ord" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" +dependencies = [ + "bitflags", +] + +[[package]] +name = "arrow-select" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + +[[package]] +name = "arrow-string" +version = "59.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + [[package]] name = "async-trait" version = "0.1.92" @@ -31,12 +208,27 @@ dependencies = [ "syn 3.0.3", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "autocfg" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "bitflags" version = "2.13.1" @@ -82,6 +274,26 @@ dependencies = [ "windows-link", ] +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + [[package]] name = "core-foundation-sys" version = "0.8.7" @@ -122,6 +334,12 @@ version = "0.8.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "displaydoc" version = "0.2.7" @@ -240,6 +458,41 @@ dependencies = [ "windows-result", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -249,13 +502,19 @@ dependencies = [ "foldhash", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "hashlink" version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" dependencies = [ - "hashbrown", + "hashbrown 0.15.5", ] [[package]] @@ -440,12 +699,75 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libc" version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "libsqlite3-sys" version = "0.35.0" @@ -515,6 +837,34 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -522,6 +872,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -684,6 +1035,12 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + [[package]] name = "redox_syscall" version = "0.5.18" @@ -693,6 +1050,18 @@ dependencies = [ "bitflags", ] +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + [[package]] name = "regex-automata" version = "0.4.18" @@ -748,6 +1117,12 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + [[package]] name = "scoped-tls" version = "1.0.1" @@ -896,6 +1271,15 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" version = "0.8.4" @@ -997,6 +1381,27 @@ version = "0.2.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + [[package]] name = "wasm-bindgen" version = "0.2.127" @@ -1130,6 +1535,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + [[package]] name = "writeable" version = "0.6.4" @@ -1159,6 +1570,26 @@ dependencies = [ "synstructure", ] +[[package]] +name = "zerocopy" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "556764e583adb45a9f8d413c2a147fa7e8d821e48e12b14fd560b607998b75eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2ab42fc20575779bd240faa45f94a74256f755c0fa9e89f0ede20d91d0cdfc1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zerofrom" version = "0.1.8" @@ -1320,6 +1751,7 @@ dependencies = [ name = "zudb-python" version = "0.0.1" dependencies = [ + "arrow", "pyo3", "zu", "zu-common", diff --git a/Cargo.toml b/Cargo.toml index 8395b87..c55c31c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -25,6 +25,11 @@ zu-common = { git = "https://github.com/tamnd/zu", rev = "67afd055032932eec36f4e # the platform it is building for, and a crate that turns the feature # on by default is a crate `cargo build` cannot link on its own. pyo3 = { version = "0.29", features = ["abi3-py311"] } +# Arrow, for the columns a result leaves as. Only `ffi` is asked for: +# the readers and writers are the engine's business and what is wanted +# here is the C Data Interface, which is how a result reaches pyarrow, +# pandas and polars without a Python object per cell. +arrow = { version = "59", default-features = false, features = ["ffi"] } [profile.release] lto = "thin" diff --git a/README.md b/README.md index 4f0a982..2f9aba5 100644 --- a/README.md +++ b/README.md @@ -46,15 +46,31 @@ zudb.load( ) with zudb.connect("social.zu1", read_only=True) as conn: - for a, b in conn.execute("MATCH (a:person)-[:knows]->(b:person) RETURN a.name AS a, b.name AS b"): + for a, b in conn.execute( + "MATCH (a:person)-[:knows]->(b:person) RETURN a.name AS a, b.name AS b" + ): print(a, "knows", b) ``` 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. +## 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. + +```python +result = conn.execute("MATCH (p:person) RETURN p.name AS name, p.score AS score") +result.to_arrow() # pyarrow.Table +result.to_pandas() # DataFrame with Arrow-backed dtypes +result.to_polars() # polars.DataFrame +result.record_batches() # a reader, for a result larger than memory +``` + +`Result` implements `__arrow_c_stream__`, so anything that reads the protocol reads a result directly and none of the four methods above is needed: `pyarrow.table(result)` and `polars.DataFrame(result)` both work. Batches are 65,536 rows. A column holds one type, which the values decide, and integers beside floats are the one mixture that widens rather than being refused. Nodes, rels and paths go across as structs. The copy runs with the GIL released, and on this machine 300,000 rows across three columns take 44 ms as Arrow against 67 ms as Python objects, and a single integer column takes 13.8 ms against 44.5 ms. + ## 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, and the GIL released around every statement and every load. Arrow, `register`, the stubs 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, 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, and the GIL released around every statement, every load and every copy out. `register`, the stubs and the interrupt are next, and each one lands with the tests that say it works. ## Wheels diff --git a/pyproject.toml b/pyproject.toml index 9902091..b4ff6ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,12 +30,24 @@ Engine = "https://github.com/tamnd/zu" Issues = "https://github.com/tamnd/zu-python/issues" [project.optional-dependencies] +# pyarrow 14 is where the Arrow PyCapsule interface arrived, and polars +# 1.3 is where its constructor started reading one. Below those a +# result would have to be handed over the old way, through a pointer +# address, which is a worse contract than saying no. arrow = ["pyarrow>=14"] pandas = ["pandas>=2.0", "pyarrow>=14"] -polars = ["polars>=1.0"] +polars = ["polars>=1.3"] +all = ["pyarrow>=14", "pandas>=2.0", "polars>=1.3"] [dependency-groups] -dev = ["maturin>=1.14,<2.0", "pytest>=8", "ruff>=0.9"] +dev = [ + "maturin>=1.14,<2.0", + "pytest>=8", + "ruff>=0.9", + "pyarrow>=14", + "pandas>=2.0", + "polars>=1.3", +] [tool.maturin] python-source = "python" diff --git a/src/columns.rs b/src/columns.rs new file mode 100644 index 0000000..cae9a8e --- /dev/null +++ b/src/columns.rs @@ -0,0 +1,673 @@ +//! A result as Arrow columns. +//! +//! This is the fast path out of the database and the reason the client +//! links the engine crates rather than the C ABI. A result comes back +//! from the executor as rows of engine values, and every one of them +//! that becomes a Python object costs an allocation, a type check and a +//! reference count. A result that becomes Arrow costs one buffer per +//! column and no Python objects at all, and pandas, polars and DuckDB +//! all read it without copying it again. +//! +//! What goes across is the Arrow C Data Interface, through the PyCapsule +//! protocol: `__arrow_c_stream__` hands out a capsule holding an +//! `ArrowArrayStream`, and every library that speaks Arrow knows how to +//! take one. There is no pyarrow dependency in the extension, and no +//! 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. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, BooleanArray, Date32Array, DurationNanosecondArray, Float64Array, Int64Array, + IntervalMonthDayNanoArray, ListArray, NullArray, StringArray, StructArray, + Time64NanosecondArray, TimestampNanosecondArray, UInt64Array, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{ + DataType, Field, FieldRef, Fields, IntervalMonthDayNano, IntervalUnit, Schema, TimeUnit, +}; +use arrow::error::ArrowError; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow::record_batch::{RecordBatch, RecordBatchIterator}; +use pyo3::prelude::*; +use zu_common::{DurationKind, Temporal}; +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. +const BATCH: usize = 65_536; + +/// What goes wrong here, with the GIL down and no way to raise yet. +/// +/// The two Python classes are the two mistakes: a value of the wrong +/// type in a column is a `TypeError`, and a value of the right type +/// that will not fit is a `ValueError`. Arrow's own errors are neither, +/// and are internal until one of them turns out to be reachable. +pub enum Snag { + Type(String), + Value(String), + Arrow(ArrowError), +} + +impl Snag { + /// The exception this is, once there is a GIL to raise it with. + pub fn raise(self, _py: Python<'_>) -> PyErr { + match self { + Snag::Type(detail) => pyo3::exceptions::PyTypeError::new_err(detail), + Snag::Value(detail) => pyo3::exceptions::PyValueError::new_err(detail), + // Nothing here is meant to be reachable: the types are + // decided before a buffer is filled, so an Arrow error is + // this module getting it wrong rather than the caller. + Snag::Arrow(err) => pyo3::exceptions::PyRuntimeError::new_err(format!( + "arrow could not build the result: {err}" + )), + } + } +} + +impl From for Snag { + fn from(err: ArrowError) -> Snag { + Snag::Arrow(err) + } +} + +/// 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)>), +} + +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(), + ), + } + } +} + +/// Every field is nullable, here and in the nested types, because a +/// null row of a struct column is a null in each of its children and +/// there is no other place to put it. +fn field(name: &str, data_type: DataType) -> FieldRef { + Arc::new(Field::new(name, data_type, true)) +} + +fn item(data_type: DataType) -> FieldRef { + field("item", data_type) +} + +fn node_fields() -> Fields { + Fields::from(vec![ + field("table", DataType::Utf8), + field("offset", DataType::UInt64), + ]) +} + +fn rel_fields() -> Fields { + Fields::from(vec![ + field("table", DataType::Utf8), + field("src", DataType::UInt64), + field("dst", DataType::UInt64), + field("ord", DataType::UInt64), + ]) +} + +fn path_fields() -> Fields { + Fields::from(vec![ + field( + "nodes", + DataType::List(item(DataType::Struct(node_fields()))), + ), + field("rels", DataType::List(item(DataType::Struct(rel_fields())))), + ]) +} + +/// An offset in minutes as the name Arrow keeps a timezone under. +/// +/// A fixed offset rather than a region, because a fixed offset is what +/// the value carries: the engine stores when a zoned datetime happened +/// and how far from UTC it was written, and no amount of arithmetic +/// recovers `Europe/Paris` from `+01:00`. +fn zone(offset: i16) -> String { + let sign = if offset < 0 { '-' } else { '+' }; + let minutes = offset.unsigned_abs(); + format!("{sign}{:02}:{:02}", minutes / 60, minutes % 60) +} + +/// 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. +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; + } + // 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 + .iter() + .map(|kind| build(kind, &[], names)) + .collect::, Snag>>()?; + batches.push(RecordBatch::try_new(schema.clone(), columns)?); + } + + 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" + )) + })?; + } + 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" + ))); + } + }, + // 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 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)?)); + } + Kind::Record(fields) + } + (left, right) if left == right => left, + _ => return None, + }) +} + +/// 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( + values + .iter() + .map(|value| match value { + Value::Bool(b) => Some(*b), + _ => None, + }) + .collect::(), + ), + Kind::Int => Arc::new( + values + .iter() + .map(|value| match value { + Value::Int(n) => Some(*n), + _ => None, + }) + .collect::(), + ), + Kind::Float => Arc::new( + values + .iter() + .map(|value| match value { + Value::Float(f) => Some(*f), + // Widened where the column holds both, which is + // the only place an integer reaches a float column. + Value::Int(n) => Some(*n as f64), + _ => None, + }) + .collect::(), + ), + Kind::Str => Arc::new( + values + .iter() + .map(|value| match value { + Value::Str(s) => Some(s.as_str()), + _ => None, + }) + .collect::(), + ), + Kind::Date => Arc::new( + temporals(values) + .map(|temporal| match temporal { + Some(Temporal::Date(days)) => Some(*days), + _ => None, + }) + .collect::(), + ), + Kind::Time => Arc::new( + temporals(values) + .map(|temporal| match temporal { + Some(Temporal::LocalTime(nanos)) => Some(*nanos), + _ => None, + }) + .collect::(), + ), + Kind::LocalDatetime => Arc::new( + temporals(values) + .map(|temporal| match temporal { + Some(Temporal::LocalDatetime(nanos)) => Some(*nanos), + _ => None, + }) + .collect::(), + ), + Kind::ZonedDatetime(offset) => Arc::new( + temporals(values) + .map(|temporal| match temporal { + Some(Temporal::ZonedDatetime { nanos, .. }) => Some(*nanos), + _ => None, + }) + .collect::() + .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 { + 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)) + } + _ => None, + }); + } + Arc::new(months.into_iter().collect::()) + } + Kind::DayTime => Arc::new( + temporals(values) + .map(|temporal| match temporal { + Some(Temporal::Duration(DurationKind::DayTime, nanos)) => Some(*nanos), + _ => None, + }) + .collect::(), + ), + Kind::Node => nodes(values, names)?, + Kind::Rel => rels(values, names)?, + Kind::Path => paths(values, names)?, + Kind::List(of) => { + let mut offsets = Vec::with_capacity(values.len() + 1); + let mut flat: Vec<&Value> = Vec::new(); + let mut valid = Vec::with_capacity(values.len()); + offsets.push(0i32); + for value in values { + if let Value::List(items) = value { + flat.extend(items.iter()); + valid.push(true); + } else { + valid.push(false); + } + offsets.push(flat.len() as i32); + } + Arc::new(ListArray::try_new( + item(of.data_type()), + OffsetBuffer::new(offsets.into()), + build(of, &flat, names)?, + Some(NullBuffer::from(valid)), + )?) + } + Kind::Record(fields) => { + let mut children: Vec = Vec::with_capacity(fields.len()); + for (at, (_, kind)) in fields.iter().enumerate() { + let column: Vec<&Value> = values + .iter() + .map(|value| match value { + Value::Record(held) => &held[at].1, + _ => &Value::Null, + }) + .collect(); + children.push(build(kind, &column, names)?); + } + Arc::new(StructArray::try_new( + match kind.data_type() { + DataType::Struct(fields) => fields, + _ => unreachable!("a record is a struct"), + }, + children, + Some(present(values)), + )?) + } + }) +} + +/// The temporal each value holds, or `None` for a value that is not one +/// and for a null. +fn temporals<'a>(values: &'a [&'a Value]) -> impl Iterator> { + values.iter().map(|value| match value { + Value::Temporal(temporal) => Some(temporal), + _ => None, + }) +} + +/// Which rows of a struct column are there at all. +fn present(values: &[&Value]) -> NullBuffer { + NullBuffer::from( + values + .iter() + .map(|value| !matches!(value, Value::Null)) + .collect::>(), + ) +} + +fn nodes(values: &[&Value], names: &Names) -> Result { + let table = tables( + values, + |value| match value { + Value::Node { table, .. } => Some(*table), + _ => None, + }, + |id| names.node_name(id), + ); + let offset: UInt64Array = values + .iter() + .map(|value| match value { + Value::Node { offset, .. } => Some(*offset), + _ => None, + }) + .collect(); + Ok(Arc::new(StructArray::try_new( + node_fields(), + vec![Arc::new(table), Arc::new(offset)], + Some(present(values)), + )?)) +} + +fn rels(values: &[&Value], names: &Names) -> Result { + let table = tables( + values, + |value| match value { + Value::Rel { table, .. } => Some(*table), + _ => None, + }, + |id| names.rel_name(id), + ); + let end = |pick: fn(&Value) -> Option| -> UInt64Array { + values.iter().map(|value| pick(value)).collect() + }; + Ok(Arc::new(StructArray::try_new( + rel_fields(), + vec![ + Arc::new(table), + Arc::new(end(|value| match value { + Value::Rel { src, .. } => Some(*src), + _ => None, + })), + Arc::new(end(|value| match value { + Value::Rel { dst, .. } => Some(*dst), + _ => None, + })), + Arc::new(end(|value| match value { + Value::Rel { ord, .. } => Some(*ord), + _ => None, + })), + ], + Some(present(values)), + )?)) +} + +/// The table name of every row, borrowed rather than copied. +/// +/// The catalog owns the names and a column holds as many rows as the +/// result does, so the names go in by reference and the only string +/// built here is the stand-in for a table the catalog no longer has, +/// which is one per missing table rather than one per row. +fn tables<'a>( + values: &[&Value], + id_of: impl Fn(&Value) -> Option, + name_of: impl Fn(u32) -> Option<&'a str>, +) -> StringArray { + let mut gone: HashMap = HashMap::new(); + for value in values { + if let Some(id) = id_of(value) + && name_of(id).is_none() + { + gone.entry(id).or_insert_with(|| format!("#{id}")); + } + } + values + .iter() + .map(|value| id_of(value).map(|id| name_of(id).unwrap_or_else(|| gone[&id].as_str()))) + .collect() +} + +/// A path column, as the two lists a walk is. +/// +/// A path is nodes and edges alternating, and Arrow has no type for a +/// list whose elements alternate between two structs. Two lists say the +/// same thing without a union in the middle of it: the nodes in the +/// order the walk visits them, the edges in the order it crosses them, +/// and one more node than edge. +fn paths(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)); + } + node_offsets.push(walked_nodes.len() as i32); + rel_offsets.push(walked_rels.len() as i32); + } + let nodes = ListArray::try_new( + item(DataType::Struct(node_fields())), + OffsetBuffer::new(node_offsets.into()), + nodes(&walked_nodes, names)?, + Some(present(values)), + )?; + let rels = ListArray::try_new( + item(DataType::Struct(rel_fields())), + OffsetBuffer::new(rel_offsets.into()), + rels(&walked_rels, names)?, + Some(present(values)), + )?; + Ok(Arc::new(StructArray::try_new( + path_fields(), + vec![Arc::new(nodes), Arc::new(rels)], + Some(present(values)), + )?)) +} diff --git a/src/conn.rs b/src/conn.rs index 47c6b58..6100992 100644 --- a/src/conn.rs +++ b/src/conn.rs @@ -7,17 +7,24 @@ //! to run statements at once want two connections. It is there so that //! a program which shares one by accident waits rather than corrupts. +use std::ffi::CStr; use std::path::PathBuf; use std::sync::Mutex; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyTuple}; +use pyo3::types::{PyCapsule, PyDict, PyList, PyTuple}; use zudb::query::{QueryResult, Value}; use zudb::{Config, Database}; +use crate::columns; use crate::error::{closed, to_py_err}; use crate::value::{Names, from_py, to_py}; +/// What a capsule holding an Arrow stream is called. The name is part +/// of the protocol: a consumer checks it before it reads the pointer, +/// and a capsule named anything else is not one of these. +const STREAM: &CStr = c"arrow_array_stream"; + /// One connection to one database. /// /// Statements run on it in order, one at a time. It reads the database @@ -254,6 +261,75 @@ impl Result { Ok(out) } + /// The rows as an Arrow stream, for anything that speaks Arrow. + /// + /// This is the PyCapsule interface, which is how a producer hands + /// Arrow data to a consumer without either of them importing the + /// other: what comes back is a capsule holding an + /// `ArrowArrayStream`, and `pyarrow.table(result)`, + /// `polars.from_arrow(result)` and anything else that reads Arrow + /// takes it from here. `requested_schema` is accepted and ignored, + /// which the protocol allows: a result has the types it has, and + /// casting them here would hide a conversion a caller can see. + #[pyo3(signature = (requested_schema = None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + // Released for the copy: filling Arrow buffers out of engine + // values touches no Python object, and a result worth putting + // in a DataFrame is big enough that another thread should get + // to run while it happens. + let stream = py + .detach(|| columns::stream(&self.result, &self.names)) + .map_err(|snag| snag.raise(py))?; + PyCapsule::new_with_value(py, stream, STREAM) + } + + /// The rows as a `pyarrow.Table`. + fn to_arrow<'py>(slf: PyRef<'py, Self>) -> PyResult> { + let py = slf.py(); + needed(py, "pyarrow", "arrow")?.call_method1("table", (slf,)) + } + + /// The rows as a `pandas.DataFrame`, with Arrow-backed dtypes. + /// + /// `ArrowDtype` rather than the NumPy dtypes pandas grew up with, + /// because the data is already Arrow and converting it to NumPy + /// would copy every column to lose null support on the way. It is + /// also what pandas 3 wants. + fn to_pandas<'py>(slf: PyRef<'py, Self>) -> PyResult> { + let py = slf.py(); + let pandas = needed(py, "pandas", "pandas")?; + let how = PyDict::new(py); + how.set_item("types_mapper", pandas.getattr("ArrowDtype")?)?; + Self::to_arrow(slf)?.call_method("to_pandas", (), Some(&how)) + } + + /// The rows as a `polars.DataFrame`. + /// + /// The constructor rather than `from_arrow`, because `from_arrow` + /// on a stream is what polars 2 is going to hand back a `Series` + /// for, and a result is a table however many columns it has. + fn to_polars<'py>(slf: PyRef<'py, Self>) -> PyResult> { + let py = slf.py(); + needed(py, "polars", "polars")?.call_method1("DataFrame", (slf,)) + } + + /// The rows as a `pyarrow.RecordBatchReader`, a batch at a time. + /// + /// 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. + fn record_batches<'py>(slf: PyRef<'py, Self>) -> PyResult> { + let py = slf.py(); + needed(py, "pyarrow", "arrow")? + .getattr("RecordBatchReader")? + .call_method1("from_stream", (slf,)) + } + fn __len__(&self) -> usize { self.result.rows.len() } @@ -283,6 +359,25 @@ impl Result { } } +/// A module the caller has to have installed for this call, imported. +/// +/// A missing one is reported as the install that fixes it rather than +/// as `ModuleNotFoundError: No module named 'pyarrow'`, which is true +/// and says nothing about what to do. An import that fails for any +/// other reason is raised as it is, because a broken pandas is not a +/// missing one. +fn needed<'py>(py: Python<'py>, module: &str, extra: &str) -> PyResult> { + PyModule::import(py, module).map_err(|err| { + if err.is_instance_of::(py) { + pyo3::exceptions::PyImportError::new_err(format!( + "this needs {module}, which is not installed: pip install 'zudb[{extra}]'" + )) + } else { + err + } + }) +} + /// The parameter dictionary as the engine takes it. fn bind(params: Option<&Bound<'_, PyDict>>) -> PyResult> { let Some(params) = params else { diff --git a/src/lib.rs b/src/lib.rs index 217dcae..7ce89c2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -12,6 +12,7 @@ //! that work. What it owes in return is the ABI's semantics, and the //! conformance corpus is what says whether it paid. +mod columns; mod conn; mod error; mod load; diff --git a/src/value.rs b/src/value.rs index 327999a..43a8ff1 100644 --- a/src/value.rs +++ b/src/value.rs @@ -53,19 +53,29 @@ impl Names { /// The table's name, or its id written out for a table the catalog /// no longer has. A result outlives nothing here, but a name is /// for reading and an unreadable one should still print. - fn node(&self, id: u32) -> String { - self.nodes - .get(&id) - .cloned() + pub fn node(&self, id: u32) -> String { + self.node_name(id) + .map(str::to_owned) .unwrap_or_else(|| format!("#{id}")) } - fn rel(&self, id: u32) -> String { - self.rels - .get(&id) - .cloned() + pub fn rel(&self, id: u32) -> String { + self.rel_name(id) + .map(str::to_owned) .unwrap_or_else(|| format!("#{id}")) } + + /// The same names, borrowed. A column of a hundred million nodes + /// holds a handful of distinct table names, and copying one of them + /// per row is the difference between building a string column and + /// allocating one. + pub fn node_name(&self, id: u32) -> Option<&str> { + self.nodes.get(&id).map(String::as_str) + } + + pub fn rel_name(&self, id: u32) -> Option<&str> { + self.rels.get(&id).map(String::as_str) + } } /// One node of the graph. diff --git a/tests/test_arrow.py b/tests/test_arrow.py new file mode 100644 index 0000000..48510ad --- /dev/null +++ b/tests/test_arrow.py @@ -0,0 +1,276 @@ +"""A result as Arrow columns, and as the frames built on them. + +The interface is the PyCapsule one, so most of these go through +pyarrow: it is the reference consumer, and what it reads out of the +capsule is what pandas, polars and everything else read too. The tests +that need pandas or polars say so and skip when they are not installed, +because the wheel depends on neither. +""" + +from __future__ import annotations + +import datetime +import threading +from pathlib import Path + +import pytest +import zudb + +pa = pytest.importorskip("pyarrow") + + +def test_a_result_is_a_table(loaded: zudb.Connection) -> None: + table = loaded.execute("MATCH (p:person) RETURN p.uid AS uid, p.name AS name").to_arrow() + assert table.num_rows == 3 + assert table.column_names == ["uid", "name"] + assert table.to_pylist() == [ + {"uid": 10, "name": "ada"}, + {"uid": 20, "name": "grace"}, + {"uid": 30, "name": "kay"}, + ] + + +def test_the_capsule_is_the_interface(loaded: zudb.Connection) -> None: + result = loaded.execute("MATCH (p:person) RETURN p.uid AS uid") + capsule = result.__arrow_c_stream__() + assert type(capsule).__name__ == "PyCapsule" + # Consumed through the protocol rather than through `to_arrow`, + # which is what every library that is not pyarrow does. + assert pa.table(result).num_rows == 3 + + +def test_a_requested_schema_is_accepted_and_the_result_is_what_it_is( + loaded: zudb.Connection, +) -> None: + # The protocol lets a consumer ask for a schema and lets a producer + # ignore it, so asking for an int32 column is not an error and does + # not get one either. + result = loaded.execute("MATCH (p:person) RETURN p.uid AS uid") + wanted = pa.schema([pa.field("uid", pa.int32())]) + capsule = result.__arrow_c_stream__(wanted.__arrow_c_schema__()) + reader = pa.RecordBatchReader._import_from_c_capsule(capsule) + assert reader.schema.field("uid").type == pa.int64() + assert reader.read_all().num_rows == 3 + + +@pytest.mark.parametrize( + "statement,params,arrow_type,answer", + [ + ("RETURN 1 AS v", {}, pa.int64(), 1), + ("RETURN 1.5 AS v", {}, pa.float64(), 1.5), + ("RETURN 'ada' AS v", {}, pa.string(), "ada"), + ("RETURN true AS v", {}, pa.bool_(), True), + ("RETURN null AS v", {}, pa.null(), None), + ("RETURN [1, 2] AS v", {}, pa.list_(pa.field("item", pa.int64())), [1, 2]), + ( + "RETURN $v AS v", + {"v": datetime.date(2024, 1, 2)}, + pa.date32(), + datetime.date(2024, 1, 2), + ), + ( + "RETURN $v AS v", + {"v": datetime.time(1, 2, 3)}, + pa.time64("ns"), + datetime.time(1, 2, 3), + ), + ( + "RETURN $v AS v", + {"v": datetime.datetime(2024, 1, 2, 3, 4, 5)}, + pa.timestamp("ns"), + datetime.datetime(2024, 1, 2, 3, 4, 5), + ), + ( + "RETURN $v AS v", + {"v": datetime.timedelta(hours=1)}, + pa.duration("ns"), + datetime.timedelta(hours=1), + ), + ], +) +def test_a_value_becomes_the_arrow_type_it_is( + empty: zudb.Connection, + statement: str, + params: dict, + arrow_type: object, + answer: object, +) -> None: + table = empty.execute(statement, params).to_arrow() + assert table.schema.field("v").type == arrow_type + got = table.column("v")[0].as_py() + if isinstance(answer, datetime.datetime): + # pyarrow gives back a pandas Timestamp when pandas is there + # and a datetime when it is not, and both compare equal to the + # datetime that went in. + assert got == answer + else: + assert got == answer + + +def test_a_zoned_datetime_carries_the_offset_it_was_written_with(empty: zudb.Connection) -> None: + zone = datetime.timezone(datetime.timedelta(hours=5, minutes=30)) + written = datetime.datetime(2024, 1, 2, 3, 4, 5, tzinfo=zone) + table = empty.execute("RETURN $v AS v", {"v": written}).to_arrow() + assert table.schema.field("v").type == pa.timestamp("ns", tz="+05:30") + assert table.column("v")[0].as_py() == written + + +def test_a_year_month_duration_is_an_interval_of_months(empty: zudb.Connection) -> None: + # Arrow has a year-month interval and pyarrow cannot build a Python + # array of one, so what goes across is the month-day-nano interval + # every reader implements, with the days and nanoseconds zero. + table = empty.execute("RETURN $v AS v", {"v": zudb.Duration(months=14)}).to_arrow() + assert table.schema.field("v").type == pa.month_day_nano_interval() + assert table.column("v")[0].as_py() == pa.MonthDayNano([14, 0, 0]) + + +def test_a_time_with_an_offset_has_no_arrow_type(empty: zudb.Connection) -> None: + zone = datetime.timezone(datetime.timedelta(hours=2)) + with pytest.raises(TypeError, match="time with an offset, which Arrow has no type for"): + empty.execute("RETURN $v AS v", {"v": datetime.time(1, 2, 3, tzinfo=zone)}).to_arrow() + + +def test_a_node_is_a_struct_of_table_and_offset(loaded: zudb.Connection) -> None: + table = loaded.execute("MATCH (p:person) RETURN p AS node").to_arrow() + assert table.schema.field("node").type == pa.struct( + [pa.field("table", pa.string()), pa.field("offset", pa.uint64())] + ) + assert table.column("node")[0].as_py() == {"table": "person", "offset": 0} + + +def test_a_rel_is_a_struct_of_its_ends(loaded: zudb.Connection) -> None: + table = loaded.execute("MATCH ()-[r:knows]->() RETURN r AS edge").to_arrow() + assert table.column("edge").to_pylist() == [ + {"table": "knows", "src": 0, "dst": 1, "ord": 0}, + {"table": "knows", "src": 1, "dst": 2, "ord": 1}, + ] + + +def test_a_path_is_its_nodes_and_its_rels(loaded: zudb.Connection) -> None: + table = loaded.execute( + "MATCH q = (a:person)-[:knows]->(b:person) WHERE a.uid = 10 RETURN q AS walk" + ).to_arrow() + assert table.column("walk").to_pylist() == [ + { + "nodes": [{"table": "person", "offset": 0}, {"table": "person", "offset": 1}], + "rels": [{"table": "knows", "src": 0, "dst": 1, "ord": 0}], + } + ] + + +def test_a_record_is_a_struct_of_its_fields(empty: zudb.Connection) -> None: + table = empty.execute("RETURN {a: 1, b: 'x'} AS rec").to_arrow() + assert table.schema.field("rec").type == pa.struct( + [pa.field("a", pa.int64()), pa.field("b", pa.string())] + ) + assert table.column("rec")[0].as_py() == {"a": 1, "b": "x"} + + +@pytest.mark.parametrize( + "statement,answer", + [ + ("UNWIND [1, null, 3] AS v RETURN v", [1, None, 3]), + ("UNWIND ['a', null] AS v RETURN v", ["a", None]), + ("UNWIND [[1, 2], null, []] AS v RETURN v", [[1, 2], None, []]), + ("UNWIND [{a: 1}, null] AS v RETURN v", [{"a": 1}, None]), + ], +) +def test_a_null_is_a_null_in_the_column_it_is_in( + empty: zudb.Connection, statement: str, answer: list +) -> None: + assert empty.execute(statement).to_arrow().column("v").to_pylist() == answer + + +def test_a_column_of_integers_and_floats_is_a_column_of_floats(empty: zudb.Connection) -> None: + table = empty.execute("UNWIND [1, 2.5] AS v RETURN v").to_arrow() + assert table.schema.field("v").type == pa.float64() + assert table.column("v").to_pylist() == [1.0, 2.5] + + +def test_a_column_of_two_types_is_refused(empty: zudb.Connection) -> None: + with pytest.raises(TypeError, match="mixes integers and strings at row 1"): + empty.execute("UNWIND [1, 'x'] AS v RETURN v").to_arrow() + + +def test_a_list_of_two_types_is_refused(empty: zudb.Connection) -> None: + with pytest.raises(TypeError, match="mixes integers and strings"): + empty.execute("RETURN [1, 'x'] AS v").to_arrow() + + +def test_a_result_with_no_rows_still_has_its_columns(loaded: zudb.Connection) -> None: + table = loaded.execute("MATCH (p:person) WHERE p.uid = 99 RETURN p.uid AS uid").to_arrow() + assert table.num_rows == 0 + assert table.column_names == ["uid"] + # Nothing said what the column holds, so it holds nothing, which + # Arrow has a type for. + assert table.schema.field("uid").type == pa.null() + + +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))}) + with zudb.connect(tmp_path / "big.zu1", read_only=True) as conn: + reader = conn.execute("MATCH (x:n) RETURN x.uid AS uid").record_batches() + batches = list(reader) + # More rows than fit in one batch, so this is two of them, and they + # add up to the result. + assert len(batches) == 2 + assert sum(batch.num_rows for batch in batches) == rows + + +def test_a_result_reads_as_arrow_and_as_objects_and_says_the_same_thing( + loaded: zudb.Connection, +) -> None: + statement = "MATCH (p:person) RETURN p.uid AS uid, p.name AS name" + objects = loaded.execute(statement).fetchall() + columns = loaded.execute(statement).to_arrow().to_pylist() + assert [tuple(row.values()) for row in columns] == objects + + +def test_a_dataframe_comes_back_with_arrow_dtypes(loaded: zudb.Connection) -> None: + pytest.importorskip("pandas") + frame = loaded.execute("MATCH (p:person) RETURN p.uid AS uid, p.name AS name").to_pandas() + assert list(frame.columns) == ["uid", "name"] + assert [str(dtype) for dtype in frame.dtypes] == ["int64[pyarrow]", "string[pyarrow]"] + assert frame["name"].tolist() == ["ada", "grace", "kay"] + + +def test_polars_reads_the_same_rows(loaded: zudb.Connection) -> None: + polars = pytest.importorskip("polars") + frame = loaded.execute("MATCH (p:person) RETURN p.uid AS uid, p.name AS name").to_polars() + assert isinstance(frame, polars.DataFrame) + assert frame.columns == ["uid", "name"] + assert frame["name"].to_list() == ["ada", "grace", "kay"] + + +def test_python_keeps_running_while_a_result_becomes_arrow(tmp_path: Path) -> None: + rows = 400_000 + zudb.load( + tmp_path / "big.zu1", + nodes="person", + rels="knows", + columns={"uid": list(range(rows)), "name": [f"p{uid}" for uid in range(rows)]}, + ) + with zudb.connect(tmp_path / "big.zu1", read_only=True) as conn: + result = conn.execute("MATCH (x:person) RETURN x.uid AS uid, x.name AS name, x AS node") + # Imported before the thread starts, so what the loop measures + # is the copy and not pyarrow's first import. + result.to_arrow() + tables: list[object] = [] + ticks = 0 + done = threading.Event() + + def run() -> None: + tables.append(result.to_arrow()) + 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() + assert tables[0].num_rows == rows + # Half a million turns on this machine, and none at all if the + # GIL were held for the copy. + assert ticks > 10_000, f"the main thread only got {ticks} turns"