From 8a45412ca92b8d47c5e71b1a4fc2f93d63414c9d Mon Sep 17 00:00:00 2001 From: 4ment Date: Sun, 20 Sep 2026 22:53:31 +1000 Subject: [PATCH 1/3] Read and write every table through the Arrow C Data Interface, and take frames from Python The store is filled from one decoder, batch_loader, which reads the vendored ArrowSchema/ArrowArray structs from their raw buffers: text in every width and as string views, lists, dates and timestamps, doubles, booleans, integers as text, the null type, and dictionary-encoded columns interned once per used value; sliced arrays and sliced batches are read at their offsets. The parquet loader exports each row group through it, so a file and a frame cannot decode differently, and a table loaded both ways is member-for-member the same store. Arrow C++ is reduced to parquet_io: a file as a C Data stream and a stream as a file. parquet_io_arrow.cpp is the only file including an Arrow header, and CPPLINK_WITH_ARROW (AUTO, ON, OFF) swaps in a stub that fails with the reason, after which the library and the command line link nothing but the C++ standard library. The merged predictions, clusters, waterfalls and the sample are built by BatchBuilder and handed to the writer as the same structs; gen-sample's output is byte-identical and the 1M load takes the same 1.7 s. Python takes a pandas frame, a dict of them, a pyarrow table or reader, or a path, all through __arrow_c_stream__, with object columns sliced on the way in. predict, rescore and cluster return DataFrames, with the reports on last_*; behind each is an ArrowTable exporting borrowed columns, the ids as a dictionary over the store's id arena, zero copy. cluster() reuses the run's rows or resolves a frame or a file through the id index. gen_sample and init work without files. The module built without Arrow installs beside PyPI's pandas and pyarrow 25 and passes the suite; CI gains a no-arrow job. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/ci.yml | 35 ++ CMakeLists.txt | 56 +- docs/python.md | 61 +- pyproject.toml | 2 +- python/bindings/common.hpp | 90 +++ python/bindings/module.cpp | 41 ++ python/bindings/schema.cpp | 62 ++ python/bindings/stages.cpp | 111 +++- python/bindings/store.cpp | 111 ++++ python/bindings/tables.cpp | 304 ++++++++++ python/cpplink/__init__.py | 400 +++++++++++-- python/tests/conftest.py | 14 +- python/tests/test_cluster.py | 15 +- python/tests/test_frames.py | 260 ++++++++ python/tests/test_linker.py | 41 +- python/tests/test_parity.py | 5 +- python/tests/test_reports.py | 32 +- src/cpplink/arrow_c.hpp | 93 +++ src/cpplink/arrow_export.cpp | 464 +++++++++++++++ src/cpplink/arrow_export.hpp | 106 ++++ src/cpplink/batch_loader.cpp | 986 +++++++++++++++++++++++++++++++ src/cpplink/batch_loader.hpp | 130 ++++ src/cpplink/cluster.cpp | 251 +++----- src/cpplink/cluster.hpp | 8 + src/cpplink/init.cpp | 42 +- src/cpplink/init.hpp | 12 + src/cpplink/merge_edges.cpp | 266 ++++++--- src/cpplink/merge_edges.hpp | 31 + src/cpplink/parquet_io.hpp | 52 ++ src/cpplink/parquet_io_arrow.cpp | 214 +++++++ src/cpplink/parquet_io_stub.cpp | 55 ++ src/cpplink/parquet_loader.cpp | 554 ++--------------- src/cpplink/parquet_loader.hpp | 12 +- src/cpplink/predict.cpp | 53 +- src/cpplink/predict.hpp | 23 + src/cpplink/rescore.cpp | 27 +- src/cpplink/rescore.hpp | 3 + src/cpplink/sample_data.cpp | 247 ++++---- src/cpplink/sample_data.hpp | 9 + src/cpplink/waterfall.cpp | 296 +++------- tests/arrow_export_test.cpp | 277 +++++++++ tests/batch_loader_test.cpp | 536 +++++++++++++++++ tests/edge_table_test.cpp | 354 +++++++++++ 43 files changed, 5447 insertions(+), 1294 deletions(-) create mode 100644 python/bindings/tables.cpp create mode 100644 python/tests/test_frames.py create mode 100644 src/cpplink/arrow_c.hpp create mode 100644 src/cpplink/arrow_export.cpp create mode 100644 src/cpplink/arrow_export.hpp create mode 100644 src/cpplink/batch_loader.cpp create mode 100644 src/cpplink/batch_loader.hpp create mode 100644 src/cpplink/parquet_io.hpp create mode 100644 src/cpplink/parquet_io_arrow.cpp create mode 100644 src/cpplink/parquet_io_stub.cpp create mode 100644 tests/arrow_export_test.cpp create mode 100644 tests/batch_loader_test.cpp create mode 100644 tests/edge_table_test.cpp diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 69d821a..31eaefb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -67,6 +67,41 @@ jobs: - name: Test run: python -m pytest python/tests + # The wheel's configuration: no conda, no Arrow, the one header-only + # dependency fetched. The library and the command line build, and the Python + # package passes its suite with PyPI's pandas and pyarrow, reading and writing + # parquet through them. + no-arrow: + name: without Arrow on ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] + defaults: + run: + shell: bash + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Configure and build the command line + run: | + cmake -S . -B build-noarrow -DCPPLINK_WITH_ARROW=OFF + cmake --build build-noarrow -j "$(nproc 2>/dev/null || sysctl -n hw.ncpu)" + ./build-noarrow/cpplink --version + + - name: Install the package + run: pip install . -Ccmake.define.CPPLINK_WITH_ARROW=OFF + + - name: Test + run: | + pip install pytest pandas pyarrow + cd "$RUNNER_TEMP" && python -m pytest "$GITHUB_WORKSPACE/python/tests" + lint: name: format and lint runs-on: ubuntu-latest diff --git a/CMakeLists.txt b/CMakeLists.txt index 60930e7..a9cec5c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -6,6 +6,15 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON) option(BUILD_TESTING "Enable building of tests" OFF) option(CPPLINK_BUILD_PYTHON "Build the Python extension module" OFF) +# Arrow C++ is what reads and writes parquet, and nothing else: every column +# crosses into the core through the C Data Interface, whose structs are vendored. +# Off, the parquet entry points fail with the reason and the library links no +# Arrow, which is how the Python wheel is built -- its input and output are data +# frames. AUTO takes Arrow where find_package finds it, which is the conda +# environment, and does without it where it does not, which is a wheel build. +# The tests need parquet fixtures, so they need it on. +set(CPPLINK_WITH_ARROW AUTO CACHE STRING "Read and write parquet through Arrow C++: ON, OFF or AUTO") +set_property(CACHE CPPLINK_WITH_ARROW PROPERTY STRINGS AUTO ON OFF) # Without this the default build has no -O flag at all, which costs roughly two # orders of magnitude in the pair-evaluation loop. Debug builds must be asked for. @@ -27,23 +36,51 @@ if(CPPLINK_SANITIZE) endif() endif() -find_package(Arrow REQUIRED) -find_package(Parquet REQUIRED) -find_package(nlohmann_json REQUIRED) +# nlohmann_json is header-only: the environment's where it has one, else the +# pinned release, fetched, which is what a wheel build on a bare runner does. +find_package(nlohmann_json QUIET) +if(NOT nlohmann_json_FOUND) + include(FetchContent) + set(JSON_BuildTests OFF CACHE INTERNAL "") + FetchContent_Declare(nlohmann_json + URL https://github.com/nlohmann/json/releases/download/v3.12.0/json.tar.xz + URL_HASH SHA256=42f6e95cad6ec532fd372391373363b62a14af6d771056dbfc86160e6dfff7aa) + FetchContent_MakeAvailable(nlohmann_json) + message(STATUS "nlohmann_json: fetched v3.12.0") +endif() + +if(CPPLINK_WITH_ARROW STREQUAL "AUTO") + find_package(Arrow QUIET) + find_package(Parquet QUIET) + if(Arrow_FOUND AND Parquet_FOUND) + set(CPPLINK_WITH_ARROW ON) + else() + set(CPPLINK_WITH_ARROW OFF) + endif() + message(STATUS "Arrow: ${CPPLINK_WITH_ARROW} (found automatically)") +elseif(CPPLINK_WITH_ARROW) + find_package(Arrow REQUIRED) + find_package(Parquet REQUIRED) +endif() # Core logic lives in a library so both the CLI and the tests can link against it. file(GLOB_RECURSE CPPLINK_SOURCES CONFIGURE_DEPENDS src/cpplink/*.cpp src/cpplink/*.hpp) +# `parquet_io.hpp` has two implementations and a build compiles one of them. +if(CPPLINK_WITH_ARROW) + list(FILTER CPPLINK_SOURCES EXCLUDE REGEX "parquet_io_stub\\.cpp$") +else() + list(FILTER CPPLINK_SOURCES EXCLUDE REGEX "parquet_io_arrow\\.cpp$") +endif() add_library(cpplink_core ${CPPLINK_SOURCES}) target_include_directories(cpplink_core PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src) # The Python extension is a shared module linking this static library, so its # objects must be relocatable. set_target_properties(cpplink_core PROPERTIES POSITION_INDEPENDENT_CODE ON) -target_link_libraries(cpplink_core PUBLIC - Arrow::arrow_shared - Parquet::parquet_shared - nlohmann_json::nlohmann_json -) +target_link_libraries(cpplink_core PUBLIC nlohmann_json::nlohmann_json) +if(CPPLINK_WITH_ARROW) + target_link_libraries(cpplink_core PUBLIC Arrow::arrow_shared Parquet::parquet_shared) +endif() add_executable(cpplink src/main.cpp) target_link_libraries(cpplink PRIVATE cpplink_core) @@ -53,6 +90,9 @@ if(CPPLINK_BUILD_PYTHON) endif() if(BUILD_TESTING) + if(NOT CPPLINK_WITH_ARROW) + message(FATAL_ERROR "the tests write parquet fixtures: BUILD_TESTING needs CPPLINK_WITH_ARROW=ON") + endif() enable_testing() find_package(GTest REQUIRED) diff --git a/docs/python.md b/docs/python.md index 848d1da..df1e1a9 100644 --- a/docs/python.md +++ b/docs/python.md @@ -3,7 +3,7 @@ The `cpplink` Python package is the same binary, loaded into the interpreter. It does two things. `cpplink.run(args)` forwards an argument list to the command line dispatcher, so every example in these pages runs unchanged as `python -m cpplink ...`. -And `cpplink.Linker` loads the parquet input once and runs each stage against the same store, returning the reports as objects rather than text. +And `cpplink.Linker` loads the input once, a pandas frame or a parquet file, and runs each stage against the same store, returning predictions and clusters as frames and the reports as objects rather than text. One rule scopes the package: no number is computed in Python. The binding calls the functions the command line calls, through the same pipeline glue, so a model estimated from Python with a seed is byte for byte the model `cpplink estimate --seed` writes. @@ -13,8 +13,7 @@ A worked example over a real dataset, with plots and tables at every stage, is t ## Install -The extension links the Arrow C++ the environment holds, which is also the Arrow that `pyarrow` loads in the same process, so the two must be the same build. -Inside the project's conda environment they are, and that is where the package is meant to be built and used: +Inside the project's conda environment: ```sh conda env update -f environment.yml @@ -26,9 +25,11 @@ python -c "import cpplink; print(cpplink.__version__)" The build goes through [scikit-build-core](https://scikit-build-core.readthedocs.io/) and lands in `build/python/`, beside the command line's `build/`. The editable install redirects the Python sources to `python/cpplink/`, so a change there needs no rebuild; a change to the bindings or the core does, with the same `pip install` line. -!!! note "No wheels" - A distributable wheel would have to bundle an Arrow matching the ABI of whatever `pyarrow` the user has, which is not a problem this package solves. - It builds for the environment it is built in. +!!! note "What the module links" + Nothing but the C++ standard library. + Every frame, table and file crosses between Python and the core through the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html), three C structs with a frozen ABI that the core vendors, so no Arrow C++ object and no `pyarrow` version is part of the contract. + `pyarrow` is what pandas uses to export a frame and what reads a parquet file into the module, and any version of it will do. + The conda build above links Arrow C++ as well, so that `predict(out="x.parquet")` and the command line write parquet themselves; a wheel built with `-DCPPLINK_WITH_ARROW=OFF` leaves that to pandas. ## The command line, in process @@ -50,20 +51,32 @@ Every method mirrors one command, with the command's flags as keyword arguments, ### 1. Make some data ```python -cpplink.gen_sample("sample.parquet", rows=200000, truth="sample.truth.csv") +frame = cpplink.gen_sample(rows=200000, truth="sample.truth.csv") # a pandas frame +cpplink.gen_sample("sample.parquet", rows=200000) # or the file ``` ### 2. Draft a schema and load ```python -schema, draft = cpplink.init("sample.parquet", out="schema.json") +schema, draft = cpplink.init(frame, out="schema.json") print(draft) # what was guessed, column by column -linker = cpplink.Linker(schema, "sample.parquet") +linker = cpplink.Linker(schema, frame) print(linker.inspect()) # cardinality and memory per column ``` -`Linker` takes a `Schema` or the path of one, and one parquet file or several. +`Linker` takes a `Schema` or the path of one, and the data: a pandas frame, a pyarrow table or reader, a polars frame or anything else with `__arrow_c_stream__`, or a parquet path. +Several inputs are a `{name: input}` dict, each becoming a dataset named by its key, or a list of paths named by their stems; two inputs and no `mode` means linking them. + +```python +linker = cpplink.Linker(schema, {"left": customers, "right": orders}) +``` + +A frame is read where it sits. +Its Arrow-backed and numeric columns cross into the core with no copy; a column of Python objects (plain strings, or a list column) is converted `batch_rows` rows at a time, so the transient is one slice and not a second frame. +A `category` column is the store's own shape and costs one intern per distinct value. +The store is a re-encoding, not a view: the frame and the store are both resident while the `Linker` lives, so drop the frame when it is not needed again. +The frame's index is not read; the record id must be a column, and an integer id is read as its digits. Several files are several datasets and default to linking; `mode="dedup"` or `"link-and-dedup"` scores the within-file pairs too, exactly as `--mode` does. `all_pairs=True` is `--all-pairs`. @@ -107,31 +120,42 @@ An `EstimateOptions` can be built and passed as `options=` instead. ### 6. Score the candidates ```python -predict = linker.predict(model, "predictions.parquet", threshold=20) -predict.predictions, predict.enumerated, predict.skipped +predictions = linker.predict(model, threshold=20) # a pandas frame +report = linker.last_predict +report.predictions, report.enumerated, report.skipped ``` -`out` is a shard directory, or a single `.csv` or `.parquet` file the shards are merged into, and the extension decides. +The frame has the merged file's columns, `id_a, id_b, gamma, match_weight, match_probability`, with `dataset_a` and `dataset_b` beside the ids over several inputs. +`out=` also writes the run as a shard directory or a single `.csv` or `.parquet` file, which `rescore` and the command line read; the extension decides. Give `threshold` in bits or `probability` as a posterior. The `--no-*` switches are the positive keywords `bounds`, `ceiling`, `signatures`, `ladders` and `interactions`; `fuzzy_tf`, `spill`, `spill_sample`, `threads`, `limit` and `tf_damping` are what they are on the command line. `verbose=True` prints the plan and writes the progress line to the process's standard error, which a terminal shows and a notebook does not capture. -The predictions are read back through pyarrow, since that file is the pipeline's door out: +`report.table` is the same rows as the core holds them, readable by anything that speaks the Arrow PyCapsule protocol and copied by none of it: ```python -import pyarrow.parquet as pq -table = pq.read_table("predictions.parquet") +import pyarrow as pa +table = pa.table(report.table) # zero copy; the ids are a dictionary over the store's ids ``` +The frame decodes that dictionary into plain string columns, since a pandas `Categorical` over every record of the store is neither cheap nor what a frame of predictions wants; the decode copies only the ids the rows name. + ### 7. Join the predictions into clusters ```python -result = linker.cluster("predictions.parquet", truth="sample.truth.csv", out="clusters.csv") +clusters = linker.cluster(truth="sample.truth.csv") # a pandas frame +result = linker.last_cluster result.quality.f1 # None without a truth file result.assignment.root # numpy uint32, the representative row per row result.assignment.size # by representative row; 0 for a non-root ``` +The frame has the cluster file's columns, `unique_id, cluster_id, cluster_size`, one row per record in a cluster of `min_size` or more, with `dataset` first over several inputs. +With nothing given, `cluster` reads the rows the last `predict` left in memory, which name rows and need no id lookup. +It also takes the frame `predict` returned, or any table with `id_a`, `id_b` and `match_weight` columns, or the shard directory or merged file a run wrote; those name records by id and cost the index the command pays. +`out=` writes the frame as `.csv` or `.parquet` too. +Which member names a cluster depends on the order the edges arrived, so two runs agree on the partition and not always on the `cluster_id`. + The numpy arrays view the C++ vectors and keep the assignment alive for as long as they do. A `Linker` holds every column, so its `cluster` costs no second load but more memory than the command, which loads the id column alone; `cpplink.cluster_file(schema_path, files, predictions)` is that lighter path. @@ -170,6 +194,5 @@ Errors the core reports as `false` and a message become `cpplink.Error`, a `Runt ## Not in this version -- **In-memory input.** The `Linker` reads parquet files, as the command does. A pyarrow, polars or pandas table through `__arrow_c_stream__` is the next step, and needs the loader split so a table can be appended without a file. -- **In-memory predictions.** `predict` writes a file and the caller reads it back; a per-thread sink handing a `pyarrow.Table` straight out is the same step. - **A progress callback.** `verbose=True` writes lines to the C-level standard error. +- **Streaming output.** `predict` holds every prediction in memory, twenty bytes each, before the frame is made; a run whose predictions do not fit writes shards with `out=` instead. diff --git a/pyproject.toml b/pyproject.toml index 67e2b05..8fa4abd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ readme = "README.md" license = {file = "LICENSE"} authors = [{name = "Mathieu Fourment"}] requires-python = ">=3.10" -dependencies = ["numpy", "pyarrow"] +dependencies = ["numpy", "pandas", "pyarrow"] [project.optional-dependencies] test = ["pytest"] diff --git a/python/bindings/common.hpp b/python/bindings/common.hpp index be2fcb0..024dedb 100644 --- a/python/bindings/common.hpp +++ b/python/bindings/common.hpp @@ -12,9 +12,12 @@ #include +#include "cpplink/arrow_export.hpp" #include "cpplink/blocking.hpp" +#include "cpplink/cluster.hpp" #include "cpplink/comparison.hpp" #include "cpplink/parquet_loader.hpp" +#include "cpplink/predict.hpp" #include "cpplink/record_store.hpp" #include "cpplink/schema.hpp" @@ -63,6 +66,14 @@ class Session { public: Session(const PySchema& schema, std::vector paths, PairMode mode, bool all_pairs); + // The same store from objects that speak the Arrow C stream protocol -- a + // pyarrow table or reader, a pandas or polars frame -- each named for the + // dataset it becomes. One input gets no name, as one file gets none. The + // frames are read through the same decoder a parquet file is, and nothing + // here links Arrow: the capsule hands over three C structs and a pointer. + Session(const PySchema& schema, + std::vector> inputs, PairMode mode, + bool all_pairs); const Schema& schema() const { return schema_; } const RecordStore& store() const { return *store_; } @@ -82,7 +93,15 @@ class Session { // --no-signatures` and `--no-ladders`. std::unique_ptr BindComparisons(bool use_signatures, bool use_ladders); + // The predictions of the last `predict`, which `cluster` with no table + // given reads: the run's own rows, nothing to resolve. + const std::shared_ptr& last_edges() const { return last_edges_; } + void set_last_edges(std::shared_ptr edges) { + last_edges_ = std::move(edges); + } + private: + std::shared_ptr last_edges_; Schema schema_; std::vector paths_; PairMode mode_; @@ -106,11 +125,82 @@ PairMode ModeFrom(const py::object& mode, size_t inputs); double ThresholdFrom(const py::object& threshold, const py::object& probability, const char* command); +// A table handed to Python through the Arrow PyCapsule protocol: whatever +// `pa.table(x)`, `pl.DataFrame(x)` or `pd.DataFrame(pa.table(x).to_pandas())` +// wants, from one builder of borrowed columns that is exported on every call and +// copies nothing. The record ids leave as a dictionary over the store's id +// arena with the rows as indices, so a table of fifty million predictions +// carries no id text of its own. `payload` is what the columns point into, and +// every live export holds it, and through it the session, until released. +class ArrowTable : public std::enable_shared_from_this { + public: + virtual ~ArrowTable() = default; + int64_t Rows() const { return builder_.Rows(); } + const std::vector& Columns() const { return names_; } + // The C structs, for a stream to hand out. Every column is borrowed, so the + // batch can be exported any number of times. + void ExportSchema(ArrowSchema* out) const { builder_.ExportSchema(out); } + bool ExportBatch(ArrowArray* out, std::string* error) { + return builder_.ExportBatch(out, error); + } + // `__arrow_c_schema__`, `__arrow_c_array__` and `__arrow_c_stream__`. + py::object SchemaCapsule() const; + py::tuple ArrayCapsules(); + py::object StreamCapsule(); + + protected: + int Add(const std::string& name, BorrowedColumn column) { + names_.push_back(name); + return builder_.AddBorrowed(name, column); + } + + BatchBuilder builder_; + std::vector names_; +}; + +// The predictions of one run: `dataset_a, id_a, dataset_b, id_b` (the datasets +// only over several inputs), `gamma`, `match_weight`, `match_probability`, the +// merged file's own columns. +class PredictionTable : public ArrowTable { + public: + PredictionTable(py::object session, std::shared_ptr edges); + const std::shared_ptr& edges() const; + + private: + struct Payload; + std::shared_ptr payload_; +}; + +// The clusters of one partition, the records in clusters of `min_size` or more: +// `dataset` (over several inputs), `unique_id`, `cluster_id`, `cluster_size`, +// the cluster file's own columns. +class ClusterTable : public ArrowTable { + public: + ClusterTable(py::object session, const ClusterAssignment& assignment, + uint64_t min_size); + + private: + struct Payload; + std::shared_ptr payload_; +}; + +// The sample `gen_sample` makes when no file is wanted. Its columns are owned +// rather than borrowed, so the rows move out on the first read and a second +// read finds none. +class SampleTable : public ArrowTable { + public: + explicit SampleTable(BatchBuilder builder) { + names_ = builder.ColumnNames(); + builder_ = std::move(builder); + } +}; + // The stages hang off the session class, so the class is made once and handed // to the files that add methods to it. using SessionClass = py::class_; void BindModule(py::module_& m); +void BindTables(py::module_& m); void BindSchema(py::module_& m); void BindModel(py::module_& m); SessionClass BindSession(py::module_& m); diff --git a/python/bindings/module.cpp b/python/bindings/module.cpp index 80ee3a2..f25daff 100644 --- a/python/bindings/module.cpp +++ b/python/bindings/module.cpp @@ -10,6 +10,7 @@ #include "bindings/common.hpp" #include "cpplink/app.hpp" +#include "cpplink/parquet_io.hpp" #include "cpplink/pipeline.hpp" #include "cpplink/sample_data.hpp" #include "cpplink/score.hpp" @@ -63,6 +64,33 @@ std::string GenSample(const std::string& out, uint64_t rows, uint64_t seed, return out; } +// The same rows with no file written: one table per output `gen-sample` would +// have made, the truth file still written where asked. +std::vector> GenSampleTables(uint64_t rows, uint64_t seed, + double duplicate_rate, + const std::string& truth, + size_t link_files) { + SampleOptions options; + options.rows = rows; + options.seed = seed; + options.duplicate_rate = duplicate_rate; + options.truth_path = truth; + options.link_paths.assign(link_files, ""); + std::vector builders; + std::string error; + bool ok = false; + { + py::gil_scoped_release release; + ok = GenerateSampleTables(options, &builders, &error); + } + Check(ok, error); + std::vector> tables; + for (BatchBuilder& builder : builders) { + tables.push_back(std::make_shared(std::move(builder))); + } + return tables; +} + } // namespace PairMode ModeFrom(const py::object& mode, size_t inputs) { @@ -89,6 +117,12 @@ double ThresholdFrom(const py::object& threshold, const py::object& probability, void BindModule(py::module_& m) { m.attr("__version__") = std::string(CPPLINK_VERSION_STRING); + m.def("parquet_supported", &ParquetSupported, + "Whether this build reads and writes parquet itself. A build without\n" + "Arrow takes its input as data frames and reads a file through pyarrow."); + m.def("dataset_names_for", &DatasetNamesFor, py::arg("paths"), + "The dataset name each file gets: its stem, made safe for a csv field and\n" + "a qualified id, and disambiguated. One file gets none."); // `kVersion` in the command line and the CMake project version are the same // number; the module carries both so a test can say so. m.attr("cli_version") = std::string(kVersion); @@ -114,6 +148,12 @@ void BindModule(py::module_& m) { "Write a sample parquet file with planted duplicates, as `cpplink gen-sample`\n" "does. `out_b` names the later files of a link fixture."); + m.def("gen_sample_tables", &GenSampleTables, py::arg("rows") = 1000000, + py::arg("seed") = 1, py::arg("duplicate_rate") = 0.08, py::arg("truth") = "", + py::arg("link_files") = 0, + "The sample as tables rather than files: one per output, the first the\n" + "originals and each later one a link fixture's duplicates."); + m.def("weight_for_probability", &WeightForProbability, py::arg("probability"), "The match weight, in bits, a match probability corresponds to."); m.def("probability_for_weight", &ProbabilityForWeight, py::arg("weight"), @@ -128,6 +168,7 @@ PYBIND11_MODULE(_cpplink, m) { cpplink::python::BindModule(m); cpplink::python::BindSchema(m); cpplink::python::BindModel(m); + cpplink::python::BindTables(m); cpplink::python::SessionClass session = cpplink::python::BindSession(m); cpplink::python::BindStages(m, &session); cpplink::python::BindDiagnostics(m, &session); diff --git a/python/bindings/schema.cpp b/python/bindings/schema.cpp index 76f3f29..5391a8b 100644 --- a/python/bindings/schema.cpp +++ b/python/bindings/schema.cpp @@ -13,6 +13,7 @@ #include #include "bindings/common.hpp" +#include "cpplink/batch_loader.hpp" #include "cpplink/init.hpp" namespace cpplink { @@ -79,6 +80,50 @@ std::pair Init( return {schema, report}; } +// The draft from inputs described by an Arrow schema each -- what a frame has +// instead of a footer. Each object speaks `__arrow_c_schema__`. +std::pair InitFrom( + const std::vector>& inputs, + const std::string& unique_id, + const std::vector>& roles, + const std::string& out) { + std::vector described; + for (const auto& [name, object] : inputs) { + if (!py::hasattr(object, "__arrow_c_schema__")) { + throw Error("input " + name + " has no __arrow_c_schema__"); + } + py::object capsule = object.attr("__arrow_c_schema__")(); + void* pointer = PyCapsule_GetPointer(capsule.ptr(), "arrow_schema"); + if (pointer == nullptr) { + PyErr_Clear(); + throw Error("__arrow_c_schema__ did not return an arrow_schema capsule"); + } + const ArrowSchema& schema = *static_cast(pointer); + DescribedInput input; + input.name = name; + for (int64_t i = 0; i < schema.n_children; ++i) { + const ArrowSchema& field = *schema.children[i]; + FileColumn column; + column.name = field.name == nullptr ? "" : field.name; + column.arrow_type = ArrowTypeName(field); + column.readable = ColumnTypeOf(field, &column.type); + input.columns.push_back(std::move(column)); + } + described.push_back(std::move(input)); + } + DraftOptions options; + options.unique_id = unique_id; + options.roles = roles; + DraftReport report; + std::string error; + Check(DraftSchemaFrom(described, options, &report, &error), error); + PySchema schema; + schema.schema = report.schema; + schema.text = report.json; + if (!out.empty()) SaveSchema(schema, out); + return {schema, report}; +} + template std::vector Describe(const std::vector& specs) { std::vector out; @@ -112,6 +157,18 @@ void BindSchema(py::module_& m) { } return names; }) + .def_property_readonly( + "input_columns", + [](const PySchema& s) { + std::vector names; + if (!s.schema.unique_id.empty()) names.push_back(s.schema.unique_id); + for (const ColumnSpec& column : s.schema.columns) { + if (!column.IsDerived()) names.push_back(column.name); + } + return names; + }, + "The columns an input must hold: the id and every column that is not\n" + "derived, which is what a reader is asked for.") .def_property_readonly( "comparison_names", [](const PySchema& s) { @@ -170,6 +227,11 @@ void BindSchema(py::module_& m) { return CaptureText([&](std::ostream& out) { PrintDraftReport(r, out); }); }); + m.def("init_from", &InitFrom, py::arg("inputs"), py::arg("id") = "", + py::arg("roles") = std::vector>{}, + py::arg("out") = "", + "Draft a schema from `(name, object)` pairs, each object speaking\n" + "`__arrow_c_schema__`: a pyarrow schema or table."); m.def("init", &Init, py::arg("files"), py::arg("id") = "", py::arg("roles") = std::vector>{}, py::arg("out") = "", diff --git a/python/bindings/stages.cpp b/python/bindings/stages.cpp index 75accb1..76ce502 100644 --- a/python/bindings/stages.cpp +++ b/python/bindings/stages.cpp @@ -41,10 +41,12 @@ struct EstimateOutcome { struct PredictOutcome : PredictReport { std::string text; // the ball tables, the plan when verbose, then the report + std::shared_ptr table; // every prediction, in memory }; struct RescoreOutcome : RescoreReport { std::string text; + std::shared_ptr table; // every prediction, in memory }; struct ClusterOutcome { @@ -53,6 +55,7 @@ struct ClusterOutcome { bool measured = false; ClusterQuality quality; std::string text; + std::shared_ptr table; // the clusters, in memory; null from a file }; struct MergeOutcome : MergeReport { @@ -93,15 +96,21 @@ bool FormatFrom(const py::object& format, EdgeFormat* out, const char* command) return true; } -PredictOutcome PredictStage(Session& session, const Model& model, const std::string& out, - const py::object& threshold, const py::object& probability, - const py::object& format, unsigned threads, uint64_t limit, - double tf_damping, bool bounds, bool ceiling, bool fuzzy_tf, +PredictOutcome PredictStage(const py::object& self, const Model& model, + const std::string& out, const py::object& threshold, + const py::object& probability, const py::object& format, + unsigned threads, uint64_t limit, double tf_damping, + bool bounds, bool ceiling, bool fuzzy_tf, const BallOptions& ball, const std::string& spill, double spill_sample, bool signatures, bool ladders, bool interactions, bool verbose) { + Session& session = py::cast(self); PredictOptions options; options.threads = threads; + // Every prediction is kept in memory whether or not a file is written: the + // table is what the front end returns, and what `cluster` reads next. + auto edges = std::make_shared(); + options.table = edges.get(); options.max_edges = limit; options.spill_dir = spill; if (spill_sample < 0.0 || spill_sample > 1.0) { @@ -110,9 +119,13 @@ PredictOutcome PredictStage(Session& session, const Model& model, const std::str options.spill_sample = spill_sample; const bool format_given = FormatFrom(format, &options.format, "predict"); std::string error; - Check(ResolveEdgeOutput("predict", out, format_given, &options.out_dir, - &options.merge_path, &error), - error); + if (!out.empty()) { + Check(ResolveEdgeOutput("predict", out, format_given, &options.out_dir, + &options.merge_path, &error), + error); + } else if (format_given) { + throw Error("predict: format names the shard format, and no out was given"); + } ScoreOptions score; score.threshold = ThresholdFrom(threshold, probability, "predict"); score.tf_damping = tf_damping; @@ -163,23 +176,32 @@ PredictOutcome PredictStage(Session& session, const Model& model, const std::str Check(ok, error); PrintPredictReport(result, scorer, text); result.text = text.str(); + session.set_last_edges(edges); + result.table = std::make_shared(self, edges); return result; } -RescoreOutcome RescoreStage(Session& session, const Model& model, +RescoreOutcome RescoreStage(const py::object& self, const Model& model, const std::string& spill, const std::string& out, const py::object& threshold, const py::object& probability, const py::object& format, unsigned threads, uint64_t limit, double tf_damping, bool bounds) { + Session& session = py::cast(self); RescoreOptions options; options.spill_dir = spill; + auto edges = std::make_shared(); + options.table = edges.get(); options.threads = threads; options.max_edges = limit; const bool format_given = FormatFrom(format, &options.format, "rescore"); std::string error; - Check(ResolveEdgeOutput("rescore", out, format_given, &options.out_dir, - &options.merge_path, &error), - error); + if (!out.empty()) { + Check(ResolveEdgeOutput("rescore", out, format_given, &options.out_dir, + &options.merge_path, &error), + error); + } else if (format_given) { + throw Error("rescore: format names the shard format, and no out was given"); + } ScoreOptions score; score.threshold = ThresholdFrom(threshold, probability, "rescore"); score.tf_damping = tf_damping; @@ -202,16 +224,16 @@ RescoreOutcome RescoreStage(Session& session, const Model& model, Check(ok, error); result.text = CaptureText( [&](std::ostream& stream) { PrintRescoreReport(result, scorer, stream); }); + session.set_last_edges(edges); + result.table = std::make_shared(self, edges); return result; } // Clustering over a store: the Linker's, or one loaded for its ids alone. -ClusterOutcome ClusterOver(const RecordStore& store, const std::string& predictions, +ClusterOutcome ClusterOver(const RecordStore& store, ClusterOptions options, const py::object& threshold, const py::object& probability, const std::string& out, uint64_t min_size, const std::string& truth) { - ClusterOptions options; - options.edge_path = predictions; options.out_path = out; options.min_size = min_size; if (!threshold.is_none() || !probability.is_none()) { @@ -242,12 +264,43 @@ ClusterOutcome ClusterOver(const RecordStore& store, const std::string& predicti return result; } -ClusterOutcome ClusterStage(Session& session, const std::string& predictions, +// `predictions` is None for the last `predict`'s own rows, a path to a shard +// directory or a merged file, or a table -- a data frame, a pyarrow table -- +// that speaks the Arrow C stream protocol and is read as a merged file is. +ClusterOutcome ClusterStage(const py::object& self, const py::object& predictions, const py::object& threshold, const py::object& probability, const std::string& out, uint64_t min_size, const std::string& truth) { - return ClusterOver(session.store(), predictions, threshold, probability, out, - min_size, truth); + Session& session = py::cast(self); + ClusterOptions options; + py::object capsule; // keeps a stream's struct alive while it is read + if (predictions.is_none()) { + if (!session.last_edges()) { + throw Error( + "cluster: no predictions in memory; run predict first, or give " + "a table or a path"); + } + options.edges = session.last_edges().get(); + } else if (py::isinstance(predictions)) { + options.edge_path = py::cast(predictions); + } else if (py::hasattr(predictions, "__arrow_c_stream__")) { + capsule = predictions.attr("__arrow_c_stream__")(); + void* pointer = PyCapsule_GetPointer(capsule.ptr(), "arrow_array_stream"); + if (pointer == nullptr) { + PyErr_Clear(); + throw Error( + "cluster: __arrow_c_stream__ did not return an arrow_array_stream"); + } + options.stream = static_cast(pointer); + } else { + throw Error( + "cluster: predictions wants None, a path, or a table with " + "__arrow_c_stream__"); + } + ClusterOutcome result = ClusterOver(session.store(), options, threshold, probability, + out, min_size, truth); + result.table = std::make_shared(self, result.assignment, min_size); + return result; } // `cpplink cluster` as the command runs it: only the id column is loaded. @@ -264,7 +317,9 @@ ClusterOutcome ClusterFile(const std::string& schema_path, ok = LoadIdsOnly(schema_path, files, &store, &error); } Check(ok, error); - return ClusterOver(*store, predictions, threshold, probability, out, min_size, truth); + ClusterOptions options; + options.edge_path = predictions; + return ClusterOver(*store, options, threshold, probability, out, min_size, truth); } MergeOutcome MergePredictions(const std::string& shards, const std::string& out, @@ -523,6 +578,8 @@ void BindStages(py::module_& m, SessionClass* session) { .def_readonly("merged_path", &PredictReport::merged_path) .def_readonly("merge_seconds", &PredictReport::merge_seconds) .def_readonly("text", &PredictOutcome::text) + .def_readonly("table", &PredictOutcome::table, + "Every prediction of the run, as an `ArrowTable`.") .def("__repr__", [](const PredictOutcome& r) { return r.text; }); py::class_(m, "SpillManifest") @@ -549,6 +606,8 @@ void BindStages(py::module_& m, SessionClass* session) { .def_readonly("merged_path", &RescoreReport::merged_path) .def_readonly("merge_seconds", &RescoreReport::merge_seconds) .def_readonly("text", &RescoreOutcome::text) + .def_readonly("table", &RescoreOutcome::table, + "Every prediction of the replay, as an `ArrowTable`.") .def("__repr__", [](const RescoreOutcome& r) { return r.text; }); py::class_( @@ -642,6 +701,9 @@ void BindStages(py::module_& m, SessionClass* session) { }, "Against the truth file, or None where none was given.") .def_readonly("text", &ClusterOutcome::text) + .def_readonly("table", &ClusterOutcome::table, + "The clusters of `min_size` or more, as an `ArrowTable`; None\n" + "from `cluster_file`, which has no session to hold.") .def("__repr__", [](const ClusterOutcome& c) { return c.text; }); py::class_(m, "MergeReport", "What `merge_predictions` wrote.") @@ -666,7 +728,7 @@ void BindStages(py::module_& m, SessionClass* session) { .def("__repr__", [](const MergeOutcome& r) { return r.text; }); session->def("estimate", &EstimateStage, py::arg("options"), py::arg("out") = "") - .def("predict", &PredictStage, py::arg("model"), py::arg("out"), + .def("predict", &PredictStage, py::arg("model"), py::arg("out") = "", py::arg("threshold") = py::none(), py::arg("probability") = py::none(), py::arg("format") = py::none(), py::arg("threads") = 0, py::arg("limit") = 0, py::arg("tf_damping") = 1.0, py::arg("bounds") = true, @@ -675,11 +737,12 @@ void BindStages(py::module_& m, SessionClass* session) { py::arg("spill_sample") = 0.0, py::arg("signatures") = true, py::arg("ladders") = true, py::arg("interactions") = true, py::arg("verbose") = false) - .def("rescore", &RescoreStage, py::arg("model"), py::arg("spill"), py::arg("out"), - py::arg("threshold") = py::none(), py::arg("probability") = py::none(), - py::arg("format") = py::none(), py::arg("threads") = 0, py::arg("limit") = 0, - py::arg("tf_damping") = 1.0, py::arg("bounds") = true) - .def("cluster", &ClusterStage, py::arg("predictions"), + .def("rescore", &RescoreStage, py::arg("model"), py::arg("spill"), + py::arg("out") = "", py::arg("threshold") = py::none(), + py::arg("probability") = py::none(), py::arg("format") = py::none(), + py::arg("threads") = 0, py::arg("limit") = 0, py::arg("tf_damping") = 1.0, + py::arg("bounds") = true) + .def("cluster", &ClusterStage, py::arg("predictions") = py::none(), py::arg("threshold") = py::none(), py::arg("probability") = py::none(), py::arg("out") = "", py::arg("min_size") = 2, py::arg("truth") = ""); diff --git a/python/bindings/store.cpp b/python/bindings/store.cpp index 960191a..9628f1c 100644 --- a/python/bindings/store.cpp +++ b/python/bindings/store.cpp @@ -1,6 +1,7 @@ // Copyright 2026 Mathieu Fourment // SPDX-License-Identifier: MIT +#include #include #include #include @@ -10,6 +11,8 @@ #include #include "bindings/common.hpp" +#include "cpplink/arrow_c.hpp" +#include "cpplink/batch_loader.hpp" #include "cpplink/id_index.hpp" #include "cpplink/inspect.hpp" #include "cpplink/pipeline.hpp" @@ -35,6 +38,101 @@ Session::Session(const PySchema& schema, std::vector paths, PairMod Check(ok, error); } +namespace { + +// The stream behind an object's `__arrow_c_stream__`: a PyCapsule named +// "arrow_array_stream" holding a pointer to the struct. The capsule keeps the +// struct's memory alive, so it is held for as long as the stream is read; the +// loader releases the stream, after which the capsule's own destructor finds +// nothing to release. +struct CapturedStream { + py::object capsule; + ArrowArrayStream* stream = nullptr; +}; + +CapturedStream StreamOf(const py::object& input, const std::string& name) { + if (!py::hasattr(input, "__arrow_c_stream__")) { + throw Error("input " + (name.empty() ? std::string() : "\"" + name + "\" ") + + "is not a data frame, a pyarrow table or anything else with " + "__arrow_c_stream__"); + } + CapturedStream captured; + captured.capsule = input.attr("__arrow_c_stream__")(); + if (!PyCapsule_CheckExact(captured.capsule.ptr())) { + throw Error("__arrow_c_stream__ did not return a capsule"); + } + void* pointer = PyCapsule_GetPointer(captured.capsule.ptr(), "arrow_array_stream"); + if (pointer == nullptr) { + PyErr_Clear(); + throw Error( + "__arrow_c_stream__ returned a capsule that is not an arrow_array_stream"); + } + captured.stream = static_cast(pointer); + return captured; +} + +// What `DatasetNamesFor` does for a file's stem, a dict key must do for itself: +// no comma, since the name sits in a csv field, and no colon, since it sits +// before the one in a qualified id. +void CheckDatasetName(const std::string& name) { + if (name.empty()) throw Error("a dataset name is empty"); + if (name.find(',') != std::string::npos || name.find(':') != std::string::npos) { + throw Error("dataset name \"" + name + "\" holds a comma or a colon"); + } +} + +} // namespace + +Session::Session(const PySchema& schema, + std::vector> inputs, PairMode mode, + bool all_pairs) + : schema_(schema.schema), mode_(mode), all_pairs_(all_pairs) { + Check(!inputs.empty(), "a Linker needs at least one input"); + std::vector streams; + std::vector named; + std::vector names; + for (const auto& [name, input] : inputs) { + if (inputs.size() > 1) { + CheckDatasetName(name); + if (std::find(names.begin(), names.end(), name) != names.end()) { + throw Error("dataset name \"" + name + "\" is given twice"); + } + } + names.push_back(name); + streams.push_back(StreamOf(input, name)); + } + // The types the schema left to the input come from the first one, as they + // come from the first file. Asking a stream for its schema reads no row. + ArrowSchema first; + first.release = nullptr; + if (streams.front().stream->get_schema(streams.front().stream, &first) != 0) { + throw Error("cannot read the schema of the first input"); + } + std::string error; + const bool typed = ResolveColumnTypesFrom( + first, inputs.size() > 1 ? names.front() : "the input", &schema_, &error); + first.release(&first); + if (!typed) { + for (const CapturedStream& s : streams) { + if (s.stream->release != nullptr) s.stream->release(s.stream); + } + throw Error(error); + } + for (size_t i = 0; i < streams.size(); ++i) { + named.push_back( + {inputs.size() > 1 ? names[i] : std::string(), streams[i].stream}); + } + store_ = std::make_unique(schema_); + bool ok = false; + { + // Pulling a batch from a pyarrow stream is native code; a stream backed + // by a Python generator takes the GIL back itself for each pull. + py::gil_scoped_release release; + ok = LoadStreams(named, schema_, store_.get(), &stats_, &error); + } + Check(ok, error); +} + const BlockingPlan& Session::Plan(bool for_estimation) { std::unique_ptr& slot = for_estimation ? estimation_plan_ : prediction_plan_; @@ -144,6 +242,19 @@ SessionClass BindSession(py::module_& m) { }), py::arg("schema"), py::arg("files"), py::arg("mode") = py::none(), py::arg("all_pairs") = false) + .def_static( + "from_streams", + [](const PySchema& schema, + std::vector> inputs, + const py::object& mode, bool all_pairs) { + const size_t count = inputs.size(); + return std::make_unique(schema, std::move(inputs), + ModeFrom(mode, count), all_pairs); + }, + py::arg("schema"), py::arg("inputs"), py::arg("mode") = py::none(), + py::arg("all_pairs") = false, + "The store from `(name, object)` pairs, each object speaking the Arrow C\n" + "stream protocol. Use `cpplink.Linker`.") .def_property_readonly("records", [](const Session& s) { return s.store().NumRecords(); }) .def_property_readonly("datasets", diff --git a/python/bindings/tables.cpp b/python/bindings/tables.cpp new file mode 100644 index 0000000..4f404fa --- /dev/null +++ b/python/bindings/tables.cpp @@ -0,0 +1,304 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "bindings/common.hpp" +#include "cpplink/arrow_c.hpp" +#include "cpplink/id_index.hpp" +#include "cpplink/score.hpp" + +namespace cpplink { +namespace python { + +namespace { + +// Strings as an arena: 64-bit offsets over one text, which is the layout a +// `large_string` column and a dictionary's values both want. +struct Arena { + std::vector offsets{0}; + std::string text; + void Add(const std::string& value) { + text += value; + offsets.push_back(static_cast(text.size())); + } + int64_t Size() const { return static_cast(offsets.size() - 1); } +}; + +BorrowedColumn Dictionary(const std::vector& indices, const int64_t* offsets, + const char* text, int64_t size) { + BorrowedColumn column; + column.type = ExportType::kDictionary; + column.length = static_cast(indices.size()); + column.values = indices.data(); + column.dictionary_offsets = offsets; + column.dictionary_text = text; + column.dictionary_size = size; + return column; +} + +// The record ids: the store's arena as the dictionary, `rows` as the indices. +BorrowedColumn IdsOf(const RecordStore& store, const std::vector& rows) { + const IdColumn& ids = store.ids(); + return Dictionary(rows, reinterpret_cast(ids.offsets.data()), + ids.text.data(), static_cast(ids.offsets.size() - 1)); +} + +BorrowedColumn NamesOf(const Arena& names, const std::vector& which) { + return Dictionary(which, names.offsets.data(), names.text.data(), names.Size()); +} + +template +BorrowedColumn Fixed(ExportType type, const std::vector& values) { + BorrowedColumn column; + column.type = type; + column.length = static_cast(values.size()); + column.values = values.data(); + return column; +} + +BorrowedColumn Strings(const Arena& arena) { + BorrowedColumn column; + column.type = ExportType::kLargeString; + column.length = arena.Size(); + column.values = arena.offsets.data(); + column.text = arena.text.c_str(); + return column; +} + +// What every payload carries: the session, so the store's arena outlives any +// export that points into it. Dropping the last reference to a Python object +// needs the GIL, and the last export may be released from anywhere. +struct SessionHold { + py::object session; + explicit SessionHold(py::object s) : session(std::move(s)) {} + ~SessionHold() { + py::gil_scoped_acquire acquire; + session = py::object(); + } +}; + +const RecordStore& StoreOf(const py::object& session) { + return py::cast(session).store(); +} + +Arena DatasetNames(const RecordStore& store) { + Arena names; + for (size_t d = 0; d < store.NumDatasets(); ++d) names.Add(store.DatasetName(d)); + return names; +} + +// What a capsule frees: the struct's own release, then the struct. +void FreeSchemaCapsule(PyObject* capsule) { + auto* schema = + static_cast(PyCapsule_GetPointer(capsule, "arrow_schema")); + if (schema == nullptr) return; + if (schema->release != nullptr) schema->release(schema); + delete schema; +} + +void FreeArrayCapsule(PyObject* capsule) { + auto* array = static_cast(PyCapsule_GetPointer(capsule, "arrow_array")); + if (array == nullptr) return; + if (array->release != nullptr) array->release(array); + delete array; +} + +void FreeStreamCapsule(PyObject* capsule) { + auto* stream = static_cast( + PyCapsule_GetPointer(capsule, "arrow_array_stream")); + if (stream == nullptr) return; + if (stream->release != nullptr) stream->release(stream); + delete stream; +} + +// A stream of one batch, exported from the table when the consumer pulls it. +// The table lives as long as the stream does. +struct OneBatchStream { + std::shared_ptr table; + bool sent = false; + std::string error; + + static int GetSchema(ArrowArrayStream* self, ArrowSchema* out) { + static_cast(self->private_data)->table->ExportSchema(out); + return 0; + } + static int GetNext(ArrowArrayStream* self, ArrowArray* out) { + auto* state = static_cast(self->private_data); + if (state->sent) { + out->release = nullptr; + return 0; + } + state->sent = true; + return state->table->ExportBatch(out, &state->error) ? 0 : EINVAL; + } + static const char* GetLastError(ArrowArrayStream* self) { + auto* state = static_cast(self->private_data); + return state->error.empty() ? nullptr : state->error.c_str(); + } + static void Release(ArrowArrayStream* self) { + delete static_cast(self->private_data); + self->release = nullptr; + } +}; + +} // namespace + +py::object ArrowTable::SchemaCapsule() const { + auto* schema = new ArrowSchema; + ExportSchema(schema); + return py::reinterpret_steal( + PyCapsule_New(schema, "arrow_schema", FreeSchemaCapsule)); +} + +py::tuple ArrowTable::ArrayCapsules() { + auto* array = new ArrowArray; + std::string error; + if (!ExportBatch(array, &error)) { + delete array; + throw Error(error); + } + py::object schema = SchemaCapsule(); + py::object batch = py::reinterpret_steal( + PyCapsule_New(array, "arrow_array", FreeArrayCapsule)); + return py::make_tuple(schema, batch); +} + +py::object ArrowTable::StreamCapsule() { + auto* state = new OneBatchStream; + state->table = shared_from_this(); + auto* stream = new ArrowArrayStream; + stream->get_schema = OneBatchStream::GetSchema; + stream->get_next = OneBatchStream::GetNext; + stream->get_last_error = OneBatchStream::GetLastError; + stream->release = OneBatchStream::Release; + stream->private_data = state; + return py::reinterpret_steal( + PyCapsule_New(stream, "arrow_array_stream", FreeStreamCapsule)); +} + +struct PredictionTable::Payload : SessionHold { + std::shared_ptr edges; + std::vector probability; + // Over several inputs: which dataset each side's record came from. + Arena names; + std::vector dataset_a; + std::vector dataset_b; + Payload(py::object s, std::shared_ptr e) + : SessionHold(std::move(s)), edges(std::move(e)) {} +}; + +PredictionTable::PredictionTable(py::object session, std::shared_ptr edges) + : payload_(std::make_shared(session, std::move(edges))) { + const RecordStore& store = StoreOf(session); + const EdgeTable& table = *payload_->edges; + const bool datasets = store.NumDatasets() > 1; + payload_->probability.reserve(table.Size()); + for (const double weight : table.weight) { + payload_->probability.push_back(ProbabilityForWeight(weight)); + } + if (datasets) { + payload_->names = DatasetNames(store); + payload_->dataset_a.reserve(table.Size()); + payload_->dataset_b.reserve(table.Size()); + for (uint64_t i = 0; i < table.Size(); ++i) { + payload_->dataset_a.push_back(store.DatasetOf(table.a[i])); + payload_->dataset_b.push_back(store.DatasetOf(table.b[i])); + } + } + // The merged file's columns, in its order. + if (datasets) Add("dataset_a", NamesOf(payload_->names, payload_->dataset_a)); + Add("id_a", IdsOf(store, table.a)); + if (datasets) Add("dataset_b", NamesOf(payload_->names, payload_->dataset_b)); + Add("id_b", IdsOf(store, table.b)); + Add("gamma", Fixed(ExportType::kUInt32, table.gamma)); + Add("match_weight", Fixed(ExportType::kDouble, table.weight)); + Add("match_probability", Fixed(ExportType::kDouble, payload_->probability)); + builder_.KeepAlive(payload_); +} + +const std::shared_ptr& PredictionTable::edges() const { + return payload_->edges; +} + +struct ClusterTable::Payload : SessionHold { + std::vector rows; // the records kept, in row order + std::vector roots; // each one's representative row + std::vector sizes; + std::vector dataset; + Arena names; + Arena qualified; // `dataset:id` per kept row, over several inputs + explicit Payload(py::object s) : SessionHold(std::move(s)) {} +}; + +ClusterTable::ClusterTable(py::object session, const ClusterAssignment& assignment, + uint64_t min_size) + : payload_(std::make_shared(session)) { + const RecordStore& store = StoreOf(session); + const bool datasets = store.NumDatasets() > 1; + Payload& p = *payload_; + for (uint64_t row = 0; row < assignment.root.size(); ++row) { + const uint32_t root = assignment.root[static_cast(row)]; + const uint64_t size = assignment.size[root]; + if (size < min_size) continue; + p.rows.push_back(static_cast(row)); + p.roots.push_back(root); + p.sizes.push_back(size); + if (datasets) { + p.dataset.push_back(store.DatasetOf(row)); + p.qualified.Add(QualifiedId(store, root)); + } + } + if (datasets) p.names = DatasetNames(store); + // The cluster file's columns, in its order. A cluster is named by its + // representative's id, so over one input that column is the arena again. + if (datasets) Add("dataset", NamesOf(p.names, p.dataset)); + Add("unique_id", IdsOf(store, p.rows)); + if (datasets) { + Add("cluster_id", Strings(p.qualified)); + } else { + Add("cluster_id", IdsOf(store, p.roots)); + } + Add("cluster_size", Fixed(ExportType::kUInt64, p.sizes)); + builder_.KeepAlive(payload_); +} + +void BindTables(py::module_& m) { + py::class_>( + m, "ArrowTable", + "A table the core holds, readable by anything that speaks the Arrow\n" + "PyCapsule protocol: `pyarrow.table(t)`, `polars.DataFrame(t)`, and through\n" + "pyarrow, pandas. Nothing is copied on the way out; the record ids are a\n" + "dictionary over the store's own ids.") + .def("__len__", &ArrowTable::Rows) + .def_property_readonly("columns", &ArrowTable::Columns) + .def("__arrow_c_schema__", &ArrowTable::SchemaCapsule) + .def( + "__arrow_c_array__", + [](ArrowTable& t, const py::object&) { return t.ArrayCapsules(); }, + py::arg("requested_schema") = py::none()) + .def( + "__arrow_c_stream__", + [](ArrowTable& t, const py::object&) { return t.StreamCapsule(); }, + py::arg("requested_schema") = py::none()); + py::class_>( + m, "PredictionTable", + "The predictions of one run, as the merged file's columns."); + py::class_>( + m, "ClusterTable", + "The clusters of one partition, as the cluster file's columns."); + py::class_>( + m, "SampleTable", + "A sample file's rows, before any file: read once, as the rows move out."); +} + +} // namespace python +} // namespace cpplink diff --git a/python/cpplink/__init__.py b/python/cpplink/__init__.py index fc98a5c..53b8475 100644 --- a/python/cpplink/__init__.py +++ b/python/cpplink/__init__.py @@ -16,8 +16,11 @@ from collections.abc import Iterable, Mapping from typing import Any +import numpy as np + from . import _cpplink from ._cpplink import ( # noqa: F401 (re-exported) + ArrowTable, BallOptions, BlockingReport, CliResult, @@ -25,6 +28,7 @@ ClusterQuality, ClusterReport, ClusterResult, + ClusterTable, CompletenessOptions, CompletenessReport, DraftReport, @@ -42,6 +46,7 @@ MissReport, Model, PairWaterfall, + PredictionTable, PredictReport, ProfileOptions, ProfileReport, @@ -60,12 +65,14 @@ from ._dict import to_dict __all__ = [ + "ArrowTable", "BallOptions", "BlockingReport", "ClusterAssignment", "ClusterQuality", "ClusterReport", "ClusterResult", + "ClusterTable", "CliResult", "CompletenessOptions", "CompletenessReport", @@ -86,6 +93,7 @@ "Model", "PairWaterfall", "PredictReport", + "PredictionTable", "ProfileOptions", "ProfileReport", "RecallMetrics", @@ -103,6 +111,7 @@ "probability_for_weight", "run", "to_dict", + "to_pandas", "weight_for_probability", ] @@ -119,6 +128,156 @@ def _path(path: PathLike | None) -> str: return "" if path is None else os.fspath(path) +# ---- inputs ----------------------------------------------------------------- + +# Rows a pandas frame with object columns is converted at a time. Every other +# input reaches the core as the buffers it already holds; an object column is a +# Python object per cell and has to be converted, and slicing bounds that copy +# to one slice rather than the frame. +BATCH_ROWS = 1 << 20 + +Input = Any # a frame, a pyarrow table or reader, or anything with __arrow_c_stream__ + +# What clustering reads of a prediction table; the rest stays where it is. +PREDICTION_COLUMNS = ["dataset_a", "id_a", "dataset_b", "id_b", "match_weight"] + + +def _is_pandas(obj: Any) -> bool: + return type(obj).__module__.split(".")[0] == "pandas" and hasattr(obj, "iloc") + + +def _is_path(obj: Any) -> bool: + return isinstance(obj, (str, os.PathLike)) + + +def _stream_of(data: Input, columns: list[str], batch_rows: int) -> Any: + """The object handed to the core: anything with ``__arrow_c_stream__``. + + A pandas frame is exported through pyarrow. Its Arrow-backed and numeric + columns cross with no copy; an object column (plain Python strings, or a + list column) is converted one slice of ``batch_rows`` at a time, from a + generator the core pulls, so the transient is a slice and not a frame. A + path is opened with pyarrow's parquet reader, which is how a build without + Arrow reads a file. Everything else is taken as it is. + """ + if _is_pandas(data): + import pyarrow as pa + + frame = data + present = [name for name in columns if name in frame.columns] + if present: + frame = frame[present] + # A numpy object column is a Python object per cell; pandas' own string + # dtype is not, whatever its `kind` says. + has_objects = any( + isinstance(dtype, np.dtype) and dtype.kind == "O" for dtype in frame.dtypes + ) + if not has_objects or len(frame) <= batch_rows: + return pa.Table.from_pandas(frame, preserve_index=False) + # Types are inferred over the whole frame, so a column empty in the first + # slice is still what the later ones hold. + schema = pa.Schema.from_pandas(frame, preserve_index=False) + + def slices() -> Any: + for start in range(0, len(frame), batch_rows): + # A table, not a batch: an Arrow-backed column may be chunked. + table = pa.Table.from_pandas( + frame.iloc[start : start + batch_rows], + schema=schema, + preserve_index=False, + ) + yield from table.to_batches() + + return pa.RecordBatchReader.from_batches(schema, slices()) + if _is_path(data): + import pyarrow.parquet as pq + + try: + reader = pq.ParquetFile(os.fspath(data)) + except OSError as failure: + raise Error(f"cannot open {os.fspath(data)}: {failure}") from failure + present = [name for name in columns if name in reader.schema_arrow.names] + batches = reader.iter_batches(columns=present or None) + schema = reader.schema_arrow + if present: + schema = pa_schema_subset(schema, present) + return _reader_from(schema, batches) + if hasattr(data, "__arrow_c_stream__"): + return data + if hasattr(data, "__arrow_c_array__"): + import pyarrow as pa + + return pa.RecordBatch.from_arrays(*_arrays_of(data)).to_reader() + raise TypeError( + f"{type(data).__name__} is not a pandas frame, a pyarrow table or reader, " + "a parquet path, or anything else with __arrow_c_stream__" + ) + + +def pa_schema_subset(schema: Any, names: list[str]) -> Any: + import pyarrow as pa + + return pa.schema([schema.field(name) for name in names], metadata=schema.metadata) + + +def _reader_from(schema: Any, batches: Any) -> Any: + import pyarrow as pa + + return pa.RecordBatchReader.from_batches(schema, batches) + + +def _arrays_of(data: Any) -> tuple[Any, Any]: + import pyarrow as pa + + batch = pa.record_batch(data) + return batch.columns, batch.schema.names + + +def _parquet_path(path: PathLike | None) -> bool: + return path is not None and os.fspath(path).lower().endswith((".parquet", ".pq")) + + +def _core_writes(path: PathLike | None) -> bool: + """Whether the core writes this output itself. A build without Arrow leaves + a parquet file to pandas, from the same rows; csv and shards it writes.""" + return not (_parquet_path(path) and not _cpplink.parquet_supported()) + + +def to_pandas(table: Any) -> Any: + """A :class:`pandas.DataFrame` from a table the core holds. + + The core hands the table over through the Arrow PyCapsule protocol with the + record ids as a dictionary over the store's own ids, which is what keeps a + table of fifty million predictions from carrying fifty million strings. + The frame decodes those to plain string columns, since a pandas + ``Categorical`` over every record of the store is neither cheap nor what a + frame of predictions wants; the decode copies only the ids the rows name. + ``pyarrow.table(table)`` gives the undecoded, zero-copy form. + """ + import pyarrow as pa + import pyarrow.compute as pc + + arrow = pa.table(table) + for index, field in enumerate(arrow.schema): + if pa.types.is_dictionary(field.type): + arrow = arrow.set_column( + index, field.name, pc.dictionary_decode(arrow[index]) + ) + return arrow.to_pandas() + + +def _inputs( + data: Input | Mapping[str, Input], columns: list[str], batch_rows: int +) -> list[tuple[str, Any]]: + """``(name, stream)`` per input, in order. A single input has no name.""" + if isinstance(data, Mapping): + return [ + (str(name), _stream_of(value, columns, batch_rows)) + for name, value in data.items() + ] + return [("", _stream_of(data, columns, batch_rows))] + + def _apply(options: Any, **values: Any) -> Any: """Set the given fields on a bound options object, refusing unknown names.""" for name, value in values.items(): @@ -153,19 +312,41 @@ def _model_to_dict(self: Model) -> dict[str, Any]: def gen_sample( - out: PathLike, + out: PathLike | None = None, rows: int = 1000000, seed: int = 1, duplicate_rate: float = 0.08, truth: PathLike | None = None, - out_b: Iterable[PathLike] = (), -) -> str: - """Write a sample parquet file with planted duplicates, as ``cpplink - gen-sample`` does. ``out_b`` names the later files of a link fixture, one - file per entry, and ``truth`` the csv of planted pairs.""" - return _cpplink.gen_sample( - os.fspath(out), rows, seed, duplicate_rate, _path(truth), _paths(out_b) + out_b: Iterable[PathLike] | int = (), +) -> Any: + """The sample with planted duplicates, as ``cpplink gen-sample`` makes it. + + With ``out`` the sample is written as parquet and the path returned; + ``out_b`` then names the later files of a link fixture, one per entry. + Without ``out`` the rows come back as a :class:`pandas.DataFrame`, or a + list of frames when ``out_b`` is a count of link files, and no file is + written but ``truth``, the csv of planted pairs, where it is asked for. + The same rows either way, from the same generator. + """ + if out is not None and not isinstance(out_b, int): + if _cpplink.parquet_supported(): + return _cpplink.gen_sample( + os.fspath(out), rows, seed, duplicate_rate, _path(truth), _paths(out_b) + ) + # A build without Arrow makes the tables and lets pandas write them. + paths = [os.fspath(out), *_paths(out_b)] + tables = _cpplink.gen_sample_tables( + rows, seed, duplicate_rate, _path(truth), len(paths) - 1 + ) + for path, table in zip(paths, tables, strict=True): + to_pandas(table).to_parquet(path, index=False) + return os.fspath(out) + link_files = out_b if isinstance(out_b, int) else len(_paths(out_b)) + tables = _cpplink.gen_sample_tables( + rows, seed, duplicate_rate, _path(truth), link_files ) + frames = [to_pandas(table) for table in tables] + return frames if link_files else frames[0] def merge_predictions( @@ -222,20 +403,67 @@ def cluster_file( ) +def _described(data: Any, name: str) -> tuple[str, Any]: + """``(name, object with __arrow_c_schema__)`` for one input to draft from.""" + import pyarrow as pa + + if _is_pandas(data): + return name, pa.Schema.from_pandas(data, preserve_index=False) + if hasattr(data, "__arrow_c_schema__"): + return name, data + if hasattr(data, "__arrow_c_stream__"): + return name, pa.RecordBatchReader.from_stream(data).schema + raise TypeError( + f"{type(data).__name__} is not a pandas frame, a pyarrow table or schema, " + "or anything else with __arrow_c_schema__" + ) + + def init( - files: PathLike | Iterable[PathLike], + data: PathLike | Iterable[PathLike] | Input | Mapping[str, Input] | None = None, id: str | None = None, # noqa: A002 (the flag's name) roles: Mapping[str, str] | None = None, out: PathLike | None = None, + *, + files: PathLike | Iterable[PathLike] | None = None, ) -> tuple[Schema, DraftReport]: - """Draft a schema from a parquet footer, as ``cpplink init`` does. - - ``roles`` overrides the guessed role of a column, as ``--role COLUMN=ROLE`` - does; the accepted roles are ``known_roles()``. Returns the schema and the - report saying what was guessed. + """Draft a schema from what the input holds, as ``cpplink init`` does. + + ``data`` is a parquet path or several, whose footers are read; or a + frame, a pyarrow table or schema, or a ``{name: input}`` dict of them, + whose Arrow schema says what the footer would have. ``roles`` overrides + the guessed role of a column, as ``--role COLUMN=ROLE`` does; the accepted + roles are ``known_roles()``. Returns the schema and the report saying + what was guessed. """ + if files is not None: + if data is not None: + raise TypeError("give data or files, not both") + data = files + if data is None: + raise TypeError("init needs data: a path, a frame, a table or a dict of them") role_list = [(column, role) for column, role in (roles or {}).items()] - return _cpplink.init(_paths(files), id or "", role_list, _path(out)) + if _is_path(data) or ( + not isinstance(data, Mapping) + and isinstance(data, Iterable) + and not hasattr(data, "__arrow_c_stream__") + and all(_is_path(item) for item in data) + ): + paths = _paths(data) + if _cpplink.parquet_supported(): + return _cpplink.init(paths, id or "", role_list, _path(out)) + import pyarrow.parquet as pq + + names = _cpplink.dataset_names_for(paths) or [paths[0]] + inputs = [ + (name, pq.read_schema(path)) for name, path in zip(names, paths, strict=True) + ] + return _cpplink.init_from(inputs, id or "", role_list, _path(out)) + if isinstance(data, Mapping): + inputs = [_described(value, str(name)) for name, value in data.items()] + else: + inputs = [_described(data, "the input")] + return _cpplink.init_from(inputs, id or "", role_list, _path(out)) # ---- Linker ----------------------------------------------------------------- @@ -244,12 +472,22 @@ def init( class Linker: """One loaded store, and every stage of the pipeline over it. - ``schema`` is a :class:`Schema` or the path of one; ``files`` one parquet - path or several, read in order into one store with each file a dataset. - ``mode`` is ``"dedup"``, ``"link"`` or ``"link-and-dedup"``, and defaults to - linking when there is more than one file. ``all_pairs`` replaces the + ``schema`` is a :class:`Schema` or the path of one. ``data`` is what to + load: a pandas frame, a pyarrow table or reader, a polars frame or + anything else with ``__arrow_c_stream__``, or a parquet path; several + inputs are a ``{name: input}`` dict, read in order into one store with + each a dataset named by its key, or a list of paths named by their stems. + ``mode`` is ``"dedup"``, ``"link"`` or ``"link-and-dedup"``, and defaults + to linking when there is more than one input. ``all_pairs`` replaces the schema's blocking with the source that blocks on nothing. + A frame is read where it sits: its Arrow-backed and numeric columns cross + into the core with no copy, and an object column is converted + ``batch_rows`` rows at a time. The store is a re-encoding rather than a + view, so the frame and the store are both resident while the ``Linker`` + lives; drop the frame if it is not needed again. The frame's index is + not read: the record id must be a column. + The estimation and prediction plans are built on first use, from the schema's ``"use": "estimate"`` and ``"use": "predict"`` sources respectively, which is the same split the command line keeps. @@ -258,15 +496,47 @@ class Linker: def __init__( self, schema: Schema | PathLike, - files: PathLike | Iterable[PathLike], + data: Input | Mapping[str, Input] | PathLike | Iterable[PathLike] | None = None, mode: str | None = None, all_pairs: bool = False, + *, + files: PathLike | Iterable[PathLike] | None = None, + batch_rows: int = BATCH_ROWS, ) -> None: if not isinstance(schema, Schema): schema = Schema.from_file(os.fspath(schema)) self.schema = schema - self.files = _paths(files) - self._session = _cpplink._Session(schema, self.files, mode, all_pairs) + self.last_predict: PredictReport | None = None + self.last_rescore: RescoreReport | None = None + self.last_cluster: ClusterResult | None = None + if files is not None: + if data is not None: + raise TypeError("give data or files, not both") + data = files + if data is None: + raise TypeError( + "a Linker needs data: a frame, a table, a path or a dict of them" + ) + self.files: list[str] = [] + if _is_path(data) or ( + not isinstance(data, Mapping) + and isinstance(data, Iterable) + and not hasattr(data, "__arrow_c_stream__") + and all(_is_path(item) for item in data) + ): + self.files = _paths(data) + if _cpplink.parquet_supported(): + self._session = _cpplink._Session(schema, self.files, mode, all_pairs) + return + # A build without Arrow reads the file through pyarrow, under the + # names the core would have given the files. + names = _cpplink.dataset_names_for(self.files) or [""] + if len(self.files) > 1: + data = dict(zip(names, self.files, strict=True)) + else: + data = self.files[0] + inputs = _inputs(data, schema.input_columns, batch_rows) + self._session = _cpplink._Session.from_streams(schema, inputs, mode, all_pairs) # -- what was loaded ------------------------------------------------------ @@ -371,7 +641,7 @@ def estimate( def predict( self, model: Model | PathLike, - out: PathLike, + out: PathLike | None = None, *, threshold: float | None = None, probability: float | None = None, @@ -389,21 +659,26 @@ def predict( ladders: bool = True, interactions: bool = True, verbose: bool = False, - ) -> PredictReport: - """Score the candidate pairs and write the predictions above a threshold. - - ``out`` is a shard directory, or a single ``.csv`` or ``.parquet`` file - the shards are merged into. Give ``threshold`` in bits or a - ``probability``. The other keywords are ``cpplink predict``'s flags, - with ``bounds``, ``ceiling``, ``signatures``, ``ladders`` and + ) -> Any: + """Score the candidate pairs and return the predictions above a threshold. + + Returns a :class:`pandas.DataFrame` with the merged file's columns: + ``id_a, id_b, gamma, match_weight, match_probability``, with + ``dataset_a`` and ``dataset_b`` beside the ids over several inputs. + The run's report is :attr:`last_predict`, and ``last_predict.table`` + is the same rows as an Arrow table, zero-copy. ``out`` also writes + them as a shard directory or a single ``.csv`` or ``.parquet`` file, + which a ``rescore`` or the command line needs. Give ``threshold`` in + bits or a ``probability``. The other keywords are ``cpplink predict``'s + flags, with ``bounds``, ``ceiling``, ``signatures``, ``ladders`` and ``interactions`` the positive forms of its ``--no-*`` switches. """ ball = BallOptions() if ball_budget is not None: ball.budget = ball_budget - return self._session.predict( + self.last_predict = self._session.predict( _model(model), - os.fspath(out), + _path(out) if _core_writes(out) else "", threshold=threshold, probability=probability, format=format, @@ -421,12 +696,16 @@ def predict( interactions=interactions, verbose=verbose, ) + frame = to_pandas(self.last_predict.table) + if not _core_writes(out): + frame.to_parquet(os.fspath(out), index=False) + return frame def rescore( self, model: Model | PathLike, spill: PathLike, - out: PathLike, + out: PathLike | None = None, *, threshold: float | None = None, probability: float | None = None, @@ -435,12 +714,17 @@ def rescore( limit: int = 0, tf_damping: float = 1.0, bounds: bool = True, - ) -> RescoreReport: - """Replay a spilled run under a new model, without comparing again.""" - return self._session.rescore( + ) -> Any: + """Replay a spilled run under a new model, without comparing again. + + Returns the predictions as :meth:`predict` does, a + :class:`pandas.DataFrame`, with the report on :attr:`last_rescore`; + ``out`` writes them as well. + """ + self.last_rescore = self._session.rescore( _model(model), os.fspath(spill), - os.fspath(out), + _path(out) if _core_writes(out) else "", threshold=threshold, probability=probability, format=format, @@ -449,33 +733,55 @@ def rescore( tf_damping=tf_damping, bounds=bounds, ) + frame = to_pandas(self.last_rescore.table) + if not _core_writes(out): + frame.to_parquet(os.fspath(out), index=False) + return frame def cluster( self, - predictions: PathLike, + predictions: Any = None, *, threshold: float | None = None, probability: float | None = None, out: PathLike | None = None, min_size: int = 2, truth: PathLike | None = None, - ) -> ClusterResult: + ) -> Any: """Join the predictions into clusters with a union-find. - ``predictions`` is the shard directory or the merged file ``predict`` - wrote. With ``truth`` the partition is scored against the known pairs - and ``result.quality`` is set. A ``Linker`` holds every column; to - cluster with only the id column loaded, as the command does, use - :func:`cluster_file`. + Returns a :class:`pandas.DataFrame` with the cluster file's columns: + ``unique_id, cluster_id, cluster_size`` for every record in a cluster + of ``min_size`` or more, with ``dataset`` first over several inputs. + ``predictions`` is ``None`` for the rows the last :meth:`predict` + left in memory, which need no id lookup; a frame ``predict`` returned, + or any table with ``id_a``, ``id_b`` and ``match_weight`` columns; or + the shard directory or merged file a run wrote. The partition, its + report and, with ``truth``, its quality are :attr:`last_cluster`. A + ``Linker`` holds every column; to cluster with only the id column + loaded, as the command does, use :func:`cluster_file`. """ - return self._session.cluster( - os.fspath(predictions), + if predictions is None: + source = None + elif _is_path(predictions): + source = os.fspath(predictions) + if _parquet_path(source) and not _cpplink.parquet_supported(): + # A build without Arrow reads the merged file through pyarrow. + source = _stream_of(source, PREDICTION_COLUMNS, BATCH_ROWS) + else: + source = _stream_of(predictions, PREDICTION_COLUMNS, BATCH_ROWS) + self.last_cluster = self._session.cluster( + source, threshold=threshold, probability=probability, - out=_path(out), + out=_path(out) if _core_writes(out) else "", min_size=min_size, truth=_path(truth), ) + frame = to_pandas(self.last_cluster.table) + if not _core_writes(out): + frame.to_parquet(os.fspath(out), index=False) + return frame def explain( self, diff --git a/python/tests/conftest.py b/python/tests/conftest.py index 0543abc..de44236 100644 --- a/python/tests/conftest.py +++ b/python/tests/conftest.py @@ -17,6 +17,15 @@ ROWS = 20000 SEED = 7 +# A build without Arrow takes frames in and hands frames out, and reads or +# writes a parquet file through pandas; what it cannot do is run the command +# line over a parquet file, or load a file's ids alone, since those are the +# core reading parquet itself. The tests of those paths are skipped there. +CORE_PARQUET = cpplink._cpplink.parquet_supported() +needs_core_parquet = pytest.mark.skipif( + not CORE_PARQUET, reason="this build reads and writes parquet through pandas only" +) + class Fixture: """A sample file, its truth pairs, the drafted schema and a Linker over it.""" @@ -34,9 +43,8 @@ def __init__(self, root: Path) -> None: out=self.model_path, seed=SEED ) self.predictions = root / "predictions.parquet" - self.predict_report = self.linker.predict( - self.model, self.predictions, threshold=10 - ) + self.frame = self.linker.predict(self.model, self.predictions, threshold=10) + self.predict_report = self.linker.last_predict class LinkFixture: diff --git a/python/tests/test_cluster.py b/python/tests/test_cluster.py index a9ed15c..088d80a 100644 --- a/python/tests/test_cluster.py +++ b/python/tests/test_cluster.py @@ -14,7 +14,8 @@ def test_views_are_one_entry_per_record(sample, tmp_path: Path) -> None: - result = sample.linker.cluster(sample.predictions, out=tmp_path / "c.csv") + sample.linker.cluster(sample.predictions, out=tmp_path / "c.csv") + result = sample.linker.last_cluster root = result.assignment.root size = result.assignment.size assert isinstance(root, np.ndarray) and root.dtype == np.uint32 @@ -38,7 +39,8 @@ def test_views_are_one_entry_per_record(sample, tmp_path: Path) -> None: def test_views_agree_with_the_written_clusters(sample, tmp_path: Path) -> None: out = tmp_path / "clusters.csv" - result = sample.linker.cluster(sample.predictions, out=out) + sample.linker.cluster(sample.predictions, out=out) + result = sample.linker.last_cluster with open(out, newline="") as handle: rows = list(csv.DictReader(handle)) assert len(rows) == result.report.written @@ -58,7 +60,8 @@ def test_views_agree_with_the_written_clusters(sample, tmp_path: Path) -> None: def test_view_outlives_the_result(sample) -> None: - root = sample.linker.cluster(sample.predictions).assignment.root + sample.linker.cluster(sample.predictions) + root = sample.linker.last_cluster.assignment.root import gc gc.collect() @@ -67,8 +70,10 @@ def test_view_outlives_the_result(sample) -> None: def test_threshold_reclusters_without_rescoring(sample) -> None: - low = sample.linker.cluster(sample.predictions) - high = sample.linker.cluster(sample.predictions, threshold=60) + sample.linker.cluster(sample.predictions) + low = sample.linker.last_cluster + sample.linker.cluster(sample.predictions, threshold=60) + high = sample.linker.last_cluster assert high.report.predictions_used < low.report.predictions_used assert high.report.clusters <= low.report.clusters assert high.report.predictions_read == low.report.predictions_read diff --git a/python/tests/test_frames.py b/python/tests/test_frames.py new file mode 100644 index 0000000..f2f3222 --- /dev/null +++ b/python/tests/test_frames.py @@ -0,0 +1,260 @@ +# Copyright 2026 Mathieu Fourment +# SPDX-License-Identifier: MIT +"""A frame in is the file in: one decoder, so the models agree to the byte. + +Every input reaches the core through the Arrow C Data Interface, and the +parquet loader goes through the same decoder, so a Linker over +``pd.read_parquet(f)`` must estimate the model the Linker over ``f`` does, +whatever dtypes the frame holds and however it is sliced on the way in. +""" + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from conftest import SEED + +import cpplink + + +def _model_text(linker: cpplink.Linker, path: Path) -> str: + linker.estimate(out=path, seed=SEED) + return path.read_text() + + +def test_a_frame_is_the_file(sample, tmp_path: Path) -> None: + frame = pd.read_parquet(sample.parquet) + linker = cpplink.Linker(sample.schema, frame) + assert linker.records == sample.linker.records + assert linker.files == [] + assert _model_text(linker, tmp_path / "frame.json") == sample.model_path.read_text() + + +def test_object_columns_are_sliced_on_the_way_in(sample, tmp_path: Path) -> None: + # Plain Python strings, a list column and a batch size that forces slices. + frame = pd.read_parquet(sample.parquet) + frame["first_name"] = frame["first_name"].astype(object) + linker = cpplink.Linker(sample.schema, frame, batch_rows=3000) + assert linker.stats.row_groups > 1, "the frame crossed in more than one batch" + assert _model_text(linker, tmp_path / "sliced.json") == sample.model_path.read_text() + + +def test_category_and_arrow_backed_columns(sample, tmp_path: Path) -> None: + frame = pd.read_parquet(sample.parquet) + frame["last_name"] = frame["last_name"].astype("category") + frame["email"] = frame["email"].astype(pd.ArrowDtype(pa.string())) + linker = cpplink.Linker(sample.schema, frame) + assert _model_text(linker, tmp_path / "typed.json") == sample.model_path.read_text() + + +def test_a_pyarrow_table_and_a_reader(sample, tmp_path: Path) -> None: + table = pq.read_table(sample.parquet) + assert ( + _model_text(cpplink.Linker(sample.schema, table), tmp_path / "table.json") + == sample.model_path.read_text() + ) + reader = pq.ParquetFile(sample.parquet).iter_batches(batch_size=4096) + stream = pa.RecordBatchReader.from_batches(pq.read_schema(sample.parquet), reader) + assert ( + _model_text(cpplink.Linker(sample.schema, stream), tmp_path / "reader.json") + == sample.model_path.read_text() + ) + + +def test_a_dict_of_frames_is_a_link(link, tmp_path: Path) -> None: + frames = {"a": pd.read_parquet(link.a), "b": pd.read_parquet(link.b)} + linker = cpplink.Linker(link.schema, frames) + assert linker.datasets == 2 + assert linker.mode == link.linker.mode + assert linker.stats.dataset_rows == link.linker.stats.dataset_rows + session = linker._session + last = linker.records - 1 + assert session.dataset_of(0) == "a" + assert session.dataset_of(last) == "b" + assert session.qualified_id(last) == "b:" + session.id_of(last) + assert _model_text(linker, tmp_path / "link.json") == _model_text( + link.linker, tmp_path / "files.json" + ) + + +def test_dict_keys_are_checked(sample) -> None: + frame = pd.read_parquet(sample.parquet) + with pytest.raises(cpplink.Error, match="comma or a colon"): + cpplink.Linker(sample.schema, {"a:b": frame, "c": frame}) + with pytest.raises(cpplink.Error, match="is empty"): + cpplink.Linker(sample.schema, {"": frame, "c": frame}) + + +def test_refusals_name_the_column(sample) -> None: + frame = pd.read_parquet(sample.parquet) + with pytest.raises(cpplink.Error, match='column "email" is not in the input'): + cpplink.Linker(sample.schema, frame.drop(columns=["email"])) + with pytest.raises(cpplink.Error, match="latitude: expected a floating point"): + cpplink.Linker(sample.schema, frame.astype({"latitude": str})) + with pytest.raises(TypeError, match="__arrow_c_stream__"): + cpplink.Linker(sample.schema, 42) + with pytest.raises(TypeError, match="not both"): + cpplink.Linker(sample.schema, frame, files=sample.parquet) + + +def test_a_path_can_be_read_through_pyarrow(sample, tmp_path: Path) -> None: + # The route a build without Arrow takes for a file: the same rows, the same model. + inputs = cpplink._inputs(sample.parquet, sample.schema.input_columns, 1 << 20) + session = cpplink._cpplink._Session.from_streams(sample.schema, inputs) + assert session.records == sample.linker.records + linker = cpplink.Linker.__new__(cpplink.Linker) + linker.schema, linker.files, linker._session = sample.schema, [], session + assert _model_text(linker, tmp_path / "pyarrow.json") == sample.model_path.read_text() + + +# ---- what comes out --------------------------------------------------------- + + +def _sorted_rows(frame: pd.DataFrame) -> pd.DataFrame: + keys = [c for c in ("dataset_a", "id_a", "dataset_b", "id_b") if c in frame.columns] + return frame.sort_values(keys).reset_index(drop=True) + + +def test_predict_returns_the_file_it_wrote(sample) -> None: + frame = sample.frame + assert list(frame.columns) == [ + "id_a", + "id_b", + "gamma", + "match_weight", + "match_probability", + ] + assert len(frame) == sample.predict_report.predictions + file = pd.read_parquet(sample.predictions) + a, b = _sorted_rows(frame), _sorted_rows(file) + assert a["id_a"].tolist() == b["id_a"].tolist() + assert a["id_b"].tolist() == b["id_b"].tolist() + assert a["gamma"].tolist() == b["gamma"].tolist() + assert a["match_weight"].tolist() == b["match_weight"].tolist() + # The ids are plain strings in the frame, and a dictionary over every record + # of the store in the table behind it, which is where zero copy lives. + assert frame["id_a"].dtype == pd.StringDtype("pyarrow", na_value=float("nan")) or str( + frame["id_a"].dtype + ).startswith("str") + table = pa.table(sample.predict_report.table) + assert pa.types.is_dictionary(table.schema.field("id_a").type) + assert len(table["id_a"].chunk(0).dictionary) == sample.linker.records + assert len(sample.predict_report.table) == len(frame) + + +def test_predict_without_a_file(sample) -> None: + frame = sample.linker.predict(sample.model, threshold=10) + assert sample.linker.last_predict.shards == [] + assert sample.linker.last_predict.merged_path == "" + assert _sorted_rows(frame).equals(_sorted_rows(sample.frame)) + with pytest.raises(cpplink.Error, match="no out was given"): + sample.linker.predict(sample.model, threshold=10, format="csv") + empty = sample.linker.predict(sample.model, threshold=1e9) + assert len(empty) == 0 and list(empty.columns) == list(frame.columns) + + +def test_cluster_from_memory_frame_and_file_agree(sample) -> None: + linker = sample.linker + linker.predict(sample.model, threshold=10) + from_memory = linker.cluster() + memory_report = linker.last_cluster.report + from_frame = linker.cluster(sample.frame) + frame_report = linker.last_cluster.report + from_file = linker.cluster(sample.predictions) + file_report = linker.last_cluster.report + assert list(from_memory.columns) == ["unique_id", "cluster_id", "cluster_size"] + assert memory_report.clusters == frame_report.clusters == file_report.clusters + assert memory_report.predictions_used == frame_report.predictions_used + + # Which member names a cluster depends on the order the edges arrived, so + # the partitions are compared as sets of members. + def partition(frame: pd.DataFrame) -> set[frozenset[str]]: + return {frozenset(g) for _, g in frame.groupby("cluster_id")["unique_id"]} + + for frame in (from_frame, from_file): + assert partition(frame) == partition(from_memory) + assert frame["cluster_size"].sum() == from_memory["cluster_size"].sum() + # The frame is what the cluster file holds. + assert len(from_memory) == file_report.written or file_report.written == 0 + assert from_memory["cluster_size"].min() >= 2 + assert set(from_memory["cluster_id"]) <= set(from_memory["unique_id"]) + + +def test_cluster_with_nothing_in_memory_says_so(sample) -> None: + frame = pd.read_parquet(sample.parquet) + fresh = cpplink.Linker(sample.schema, frame) + with pytest.raises(cpplink.Error, match="no predictions in memory"): + fresh.cluster() + with pytest.raises(TypeError, match="__arrow_c_stream__"): + fresh.cluster(42) + + +def test_link_frames_carry_the_datasets(link) -> None: + linker = cpplink.Linker( + link.schema, {"a": pd.read_parquet(link.a), "b": pd.read_parquet(link.b)} + ) + model, _ = linker.estimate(seed=3) + predictions = linker.predict(model, threshold=10) + assert list(predictions.columns)[:4] == ["dataset_a", "id_a", "dataset_b", "id_b"] + assert set(predictions["dataset_a"]) == {"a"} + assert set(predictions["dataset_b"]) == {"b"} + clusters = linker.cluster(truth=link.truth) + assert list(clusters.columns) == [ + "dataset", + "unique_id", + "cluster_id", + "cluster_size", + ] + assert set(clusters["dataset"]) == {"a", "b"} + assert clusters["cluster_id"].str.startswith("a:").all() + assert linker.last_cluster.quality.f1 > 0.99 + again = linker.cluster(predictions) + assert len(again) == len(clusters) + assert linker.last_cluster.report.clusters == clusters["cluster_id"].nunique() + + +# ---- fixtures without files ------------------------------------------------- + + +def test_gen_sample_frames_are_the_files_rows(sample, tmp_path: Path) -> None: + from conftest import ROWS + + frame = cpplink.gen_sample(rows=ROWS, truth=tmp_path / "truth.csv") + file = pd.read_parquet(sample.parquet) + assert len(frame) == ROWS + assert list(frame.columns) == list(file.columns) + assert frame["id"].tolist() == file["id"].tolist() + assert frame["email"].tolist() == file["email"].tolist() + assert frame["dob"].tolist() == file["dob"].tolist() + assert [list(t) for t in frame["address_tokens"]] == [ + list(t) for t in file["address_tokens"] + ] + assert (tmp_path / "truth.csv").read_text() == sample.truth.read_text() + a, b = cpplink.gen_sample(rows=ROWS, out_b=1) + assert len(a) + len(b) == ROWS + assert ( + b["id"].tolist() + == pd.read_parquet( + cpplink.gen_sample( + tmp_path / "a.parquet", rows=ROWS, out_b=[tmp_path / "b.parquet"] + ).replace("a.parquet", "b.parquet") + )["id"].tolist() + ) + + +def test_init_from_a_frame_is_init_from_the_file(sample, link) -> None: + frame = pd.read_parquet(sample.parquet) + schema, report = cpplink.init(frame) + assert schema.text == sample.schema.text + assert [c.role for c in report.columns] == [c.role for c in sample.draft.columns] + table = pq.read_table(sample.parquet) + assert cpplink.init(table)[0].text == sample.schema.text + assert cpplink.init(table.schema)[0].text == sample.schema.text + both = cpplink.init({"a": pd.read_parquet(link.a), "b": pd.read_parquet(link.b)}) + assert both[0].text == link.schema.text + with pytest.raises(cpplink.Error, match="is in a but not in b"): + cpplink.init({"a": frame, "b": frame.drop(columns=["phone"])}) diff --git a/python/tests/test_linker.py b/python/tests/test_linker.py index fc7160d..9bc3f8c 100644 --- a/python/tests/test_linker.py +++ b/python/tests/test_linker.py @@ -8,7 +8,7 @@ import pyarrow.parquet as pq import pytest -from conftest import ROWS +from conftest import CORE_PARQUET, ROWS, needs_core_parquet import cpplink @@ -73,7 +73,8 @@ def test_estimate_refuses_unknown_option(sample) -> None: def test_predict_writes_one_file(sample) -> None: report = sample.predict_report - assert report.merged_path == str(sample.predictions) + # A build without Arrow hands the frame to pandas to write instead. + assert report.merged_path == (str(sample.predictions) if CORE_PARQUET else "") assert report.predictions > 0 assert report.enumerated > report.predictions assert report.mode == "all pairs" @@ -99,9 +100,12 @@ def test_predict_needs_a_threshold(sample, tmp_path: Path) -> None: ) +@needs_core_parquet def test_predict_to_shards_and_merge(sample, tmp_path: Path) -> None: shards = tmp_path / "shards" - report = sample.linker.predict(sample.model, shards, threshold=10, threads=2) + frame = sample.linker.predict(sample.model, shards, threshold=10, threads=2) + report = sample.linker.last_predict + assert len(frame) == report.predictions assert report.merged_path == "" assert len(report.shards) == 2 merged = tmp_path / "merged.csv" @@ -119,7 +123,9 @@ def test_predict_to_shards_and_merge(sample, tmp_path: Path) -> None: def test_cluster_scores_the_planted_pairs(sample, tmp_path: Path) -> None: out = tmp_path / "clusters.csv" - result = sample.linker.cluster(sample.predictions, truth=sample.truth, out=out) + frame = sample.linker.cluster(sample.predictions, truth=sample.truth, out=out) + result = sample.linker.last_cluster + assert len(frame) == result.report.written assert result.quality is not None # Not exactly 1.0: `gen_sample` draws through the standard library's # distributions, whose output differs between libstdc++ and libc++, so the @@ -133,14 +139,17 @@ def test_cluster_scores_the_planted_pairs(sample, tmp_path: Path) -> None: assert result.report.written > 0 assert out.exists() assert "F1" in result.text or "f1" in result.text.lower() - without = sample.linker.cluster(sample.predictions) + sample.linker.cluster(sample.predictions) + without = sample.linker.last_cluster assert without.quality is None assert without.report.clusters == result.report.clusters +@needs_core_parquet def test_cluster_file_matches_the_linker(sample) -> None: light = cpplink.cluster_file(sample.schema_path, [sample.parquet], sample.predictions) - full = sample.linker.cluster(sample.predictions) + sample.linker.cluster(sample.predictions) + full = sample.linker.last_cluster assert light.report.clusters == full.report.clusters assert (light.assignment.root == full.assignment.root).all() @@ -249,13 +258,17 @@ def test_completeness(sample) -> None: def test_rescore_replays_a_spill(sample, tmp_path: Path) -> None: spill = tmp_path / "spill" - first = sample.linker.predict( + sample.linker.predict( sample.model, tmp_path / "first.parquet", threshold=10, spill=spill ) + first = sample.linker.last_predict assert first.spilled > 0 - report = sample.linker.rescore( + frame = sample.linker.rescore( sample.model, spill, tmp_path / "again.parquet", threshold=20 ) + report = sample.linker.last_rescore + assert len(frame) == report.predictions + assert frame["match_weight"].min() >= 20 assert report.pairs == first.spilled assert report.predictions <= first.predictions assert not report.below_spill_threshold @@ -263,9 +276,8 @@ def test_rescore_replays_a_spill(sample, tmp_path: Path) -> None: again = pq.read_table(str(tmp_path / "again.parquet")) assert again.num_rows == report.predictions assert min(again["match_weight"].to_pylist()) >= 20 - lower = sample.linker.rescore( - sample.model, spill, tmp_path / "lower.parquet", threshold=5 - ) + sample.linker.rescore(sample.model, spill, tmp_path / "lower.parquet", threshold=5) + lower = sample.linker.last_rescore assert lower.below_spill_threshold assert "below" in lower.text.lower() @@ -286,13 +298,16 @@ def test_link_mode(link, tmp_path: Path) -> None: assert recall.metrics.pair_completeness > 0.99 model, _ = linker.estimate(seed=3) out = tmp_path / "link.parquet" - predict = linker.predict(model, out, threshold=10) + frame = linker.predict(model, out, threshold=10) + predict = linker.last_predict assert predict.datasets == 2 + assert list(frame.columns)[:4] == ["dataset_a", "id_a", "dataset_b", "id_b"] table = pq.read_table(str(out)) assert table.schema.names[:4] == ["dataset_a", "id_a", "dataset_b", "id_b"] assert set(table["dataset_a"].to_pylist()) == {"a"} assert set(table["dataset_b"].to_pylist()) == {"b"} - result = linker.cluster(out, truth=link.truth) + linker.cluster(out, truth=link.truth) + result = linker.last_cluster assert result.quality.f1 > 0.99 row = table.slice(0, 1).to_pylist()[0] explanation = linker.explain(f"a:{row['id_a']}", f"b:{row['id_b']}", model=model) diff --git a/python/tests/test_parity.py b/python/tests/test_parity.py index db7f42b..cfedb1f 100644 --- a/python/tests/test_parity.py +++ b/python/tests/test_parity.py @@ -13,10 +13,13 @@ from pathlib import Path import pyarrow.parquet as pq -from conftest import SEED +from conftest import SEED, needs_core_parquet import cpplink +# The command line reads the parquet file itself. +pytestmark = needs_core_parquet + def _run(*args: str) -> cpplink.CliResult: result = cpplink.run([str(arg) for arg in args]) diff --git a/python/tests/test_reports.py b/python/tests/test_reports.py index b0448a5..2d672fb 100644 --- a/python/tests/test_reports.py +++ b/python/tests/test_reports.py @@ -12,6 +12,7 @@ import json import pytest +from conftest import CORE_PARQUET import cpplink from cpplink import to_dict @@ -28,10 +29,10 @@ "SessionFit": 14, "PairResidual": 9, "InteractionReport": 8, - "PredictReport": 20, - "RescoreReport": 12, + "PredictReport": 21, + "RescoreReport": 13, "SpillManifest": 8, - "ClusterResult": 3, + "ClusterResult": 4, "ClusterAssignment": 7, "ClusterReport": 19, "ClusterQuality": 9, @@ -61,20 +62,20 @@ def reports(sample, tmp_path_factory: pytest.TempPathFactory) -> dict[str, objec linker = sample.linker model = sample.model spill = root / "spill" - predict = linker.predict(model, root / "p.parquet", threshold=10, spill=spill) + linker.predict(model, root / "p.parquet", threshold=10, spill=spill) + predict = linker.last_predict shards = root / "shards" linker.predict(model, shards, threshold=10, threads=2) - cluster = linker.cluster(sample.predictions, truth=sample.truth) + linker.cluster(sample.predictions, truth=sample.truth) + cluster = linker.last_cluster explanation = linker.explain(linker.id_of(0), linker.id_of(1), model=model) recall = linker.recall(sample.truth, why=True) levels, _ = linker.levels() simplify, _ = linker.simplify(model) inspection = linker.inspect() - merge = cpplink.merge_predictions( - shards, root / "merged.csv", schema=sample.schema_path, files=[sample.parquet] - ) - rescore = linker.rescore(model, spill, root / "r.parquet", threshold=12) - return { + linker.rescore(model, spill, root / "r.parquet", threshold=12) + rescore = linker.last_rescore + reports = { "Inspection": inspection, "LoadStats": inspection.stats, "MemoryReport": inspection.memory, @@ -93,7 +94,6 @@ def reports(sample, tmp_path_factory: pytest.TempPathFactory) -> dict[str, objec "ClusterAssignment": cluster.assignment, "ClusterReport": cluster.report, "ClusterQuality": cluster.quality, - "MergeReport": merge, "Explanation": explanation, "PairWaterfall": explanation.waterfall, "WaterfallStep": explanation.waterfall.steps[0], @@ -111,10 +111,18 @@ def reports(sample, tmp_path_factory: pytest.TempPathFactory) -> dict[str, objec "CompletenessReport": linker.completeness(model), "DraftReport": sample.draft, } + if CORE_PARQUET: + # Naming binary shards' rows needs the file's ids, which the core loads. + reports["MergeReport"] = cpplink.merge_predictions( + shards, root / "merged.csv", schema=sample.schema_path, files=[sample.parquet] + ) + return reports @pytest.mark.parametrize("name", sorted(FIELDS)) def test_to_dict_sees_every_field(reports, name: str) -> None: + if name not in reports: + pytest.skip("this build reads and writes parquet through pandas only") report = reports[name] assert type(report).__name__ == name as_dict = to_dict(report) @@ -126,6 +134,8 @@ def test_to_dict_sees_every_field(reports, name: str) -> None: @pytest.mark.parametrize("name", sorted(FIELDS)) def test_repr_is_not_empty(reports, name: str) -> None: + if name not in reports: + pytest.skip("this build reads and writes parquet through pandas only") assert repr(reports[name]).strip() diff --git a/src/cpplink/arrow_c.hpp b/src/cpplink/arrow_c.hpp new file mode 100644 index 0000000..fe84365 --- /dev/null +++ b/src/cpplink/arrow_c.hpp @@ -0,0 +1,93 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#pragma once + +// The Arrow C Data Interface and C Stream Interface, as the specification defines +// them: https://arrow.apache.org/docs/format/CDataInterface.html +// +// Three plain C structs with a frozen ABI are the whole contract between cpplink +// and anything that holds Arrow-shaped memory -- pyarrow, pandas, polars, duckdb, +// or Arrow C++ itself. Vendoring them here is what lets the record store be filled +// from a data frame without linking libarrow: a consumer reads the buffers the +// producer points at, and calls `release` when it is done. The include guards are +// the specification's own, so this header and Arrow's `arrow/c/abi.h` can meet in +// one translation unit without redefining anything. + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#ifndef ARROW_C_DATA_INTERFACE +#define ARROW_C_DATA_INTERFACE + +#define ARROW_FLAG_DICTIONARY_ORDERED 1 +#define ARROW_FLAG_NULLABLE 2 +#define ARROW_FLAG_MAP_KEYS_SORTED 4 + +struct ArrowSchema { + // Array type description + const char* format; + const char* name; + const char* metadata; + int64_t flags; + int64_t n_children; + struct ArrowSchema** children; + struct ArrowSchema* dictionary; + + // Release callback + void (*release)(struct ArrowSchema*); + // Opaque producer-specific data + void* private_data; +}; + +struct ArrowArray { + // Array data description + int64_t length; + int64_t null_count; + int64_t offset; + int64_t n_buffers; + int64_t n_children; + const void** buffers; + struct ArrowArray** children; + struct ArrowArray* dictionary; + + // Release callback + void (*release)(struct ArrowArray*); + // Opaque producer-specific data + void* private_data; +}; + +#endif // ARROW_C_DATA_INTERFACE + +#ifndef ARROW_C_STREAM_INTERFACE +#define ARROW_C_STREAM_INTERFACE + +struct ArrowArrayStream { + // Callback to get the stream type (the same for every array in the stream). + // Returns 0 on success, an errno-compatible code otherwise. On success the + // ArrowSchema must be released independently of the stream. + int (*get_schema)(struct ArrowArrayStream*, struct ArrowSchema* out); + + // Callback to get the next array. A released `out` with no error means the + // stream has ended. On success the ArrowArray must be released independently + // of the stream. + int (*get_next)(struct ArrowArrayStream*, struct ArrowArray* out); + + // Optional detail on the last failure, valid until the next operation on the + // stream. Only to be called after a non-zero return. + const char* (*get_last_error)(struct ArrowArrayStream*); + + // Release callback: release the stream's own resources. + void (*release)(struct ArrowArrayStream*); + // Opaque producer-specific data + void* private_data; +}; + +#endif // ARROW_C_STREAM_INTERFACE + +#ifdef __cplusplus +} +#endif diff --git a/src/cpplink/arrow_export.cpp b/src/cpplink/arrow_export.cpp new file mode 100644 index 0000000..c813d48 --- /dev/null +++ b/src/cpplink/arrow_export.cpp @@ -0,0 +1,464 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#include "cpplink/arrow_export.hpp" + +#include +#include +#include +#include +#include +#include + +namespace cpplink { + +// The buffers of one column, in the layout the C Data Interface wants them: a +// validity bitmap where any value is null, offsets and bytes for text, packed +// bits for booleans, a child column for the elements of a list. +struct BatchBuilder::Column { + std::string name; + ExportType type = ExportType::kString; + int64_t length = 0; + int64_t null_count = 0; + std::vector validity; // grown only once a null is appended + std::vector offsets; // text and lists; starts at 0 + std::vector data; // text bytes, packed bits or fixed-width values + std::unique_ptr child; // the elements of a list + bool borrowed = false; + BorrowedColumn borrow; + + void SetValid(bool valid) { + const size_t byte = static_cast(length >> 3); + if (!valid && validity.empty()) { + // Every row so far was valid: fill the bitmap in retroactively. + validity.assign(byte + 1, 0xFF); + } + if (!validity.empty()) { + if (validity.size() <= byte) validity.push_back(0xFF); + if (valid) { + validity[byte] |= static_cast(1u << (length & 7)); + } else { + validity[byte] &= static_cast(~(1u << (length & 7))); + ++null_count; + } + } + ++length; + } + + template + void PutFixed(T value) { + const size_t at = data.size(); + data.resize(at + sizeof(T)); + std::memcpy(data.data() + at, &value, sizeof(T)); + } + + void PutBit(bool value) { + const size_t byte = static_cast(length >> 3); + if (data.size() <= byte) data.push_back(0); + if (value) data[byte] |= static_cast(1u << (length & 7)); + } + + void Clear() { + length = 0; + null_count = 0; + validity.clear(); + offsets.clear(); + data.clear(); + if (child) child->Clear(); + } +}; + +namespace { + +const char* FormatOf(ExportType type) { + switch (type) { + case ExportType::kString: + return "u"; + case ExportType::kBoolean: + return "b"; + case ExportType::kUInt8: + return "C"; + case ExportType::kUInt32: + return "I"; + case ExportType::kUInt64: + return "L"; + case ExportType::kDouble: + return "g"; + case ExportType::kDate32: + return "tdD"; + case ExportType::kStringList: + return "+l"; + case ExportType::kDictionary: + return "I"; // the indices; the values are the dictionary's own schema + case ExportType::kLargeString: + return "U"; + } + return ""; +} + +// What a released schema or array frees. Both hold their children inline, so a +// consumer releasing the parent releases the whole tree; the children's own +// release callbacks only mark them released, as the specification asks. +struct OwnedSchema { + std::deque names; // a deque keeps every c_str() where it was + std::vector children; + std::vector child_pointers; + std::vector> nested; +}; + +void ReleaseChildSchema(ArrowSchema* schema) { schema->release = nullptr; } + +void ReleaseSchema(ArrowSchema* schema) { + delete static_cast(schema->private_data); + schema->release = nullptr; +} + +// Fills `out` for one column, hanging its child (a list's elements) off it. +void FillSchema(const BatchBuilder::Column& column, OwnedSchema* owner, + ArrowSchema* out) { + owner->names.push_back(column.name); + out->format = FormatOf(column.type); + out->name = owner->names.back().c_str(); + out->metadata = nullptr; + out->flags = ARROW_FLAG_NULLABLE; + out->n_children = 0; + out->children = nullptr; + out->dictionary = nullptr; + out->release = ReleaseChildSchema; + out->private_data = nullptr; + if (column.child) { + auto nested = std::make_unique(); + nested->children.resize(1); + FillSchema(*column.child, nested.get(), &nested->children[0]); + nested->child_pointers = {&nested->children[0]}; + out->n_children = 1; + out->children = nested->child_pointers.data(); + owner->nested.push_back(std::move(nested)); + } + if (column.type == ExportType::kDictionary) { + auto nested = std::make_unique(); + nested->children.resize(1); + ArrowSchema* values = &nested->children[0]; + values->format = "U"; + values->name = ""; + values->metadata = nullptr; + values->flags = 0; + values->n_children = 0; + values->children = nullptr; + values->dictionary = nullptr; + values->release = ReleaseChildSchema; + values->private_data = nullptr; + out->dictionary = values; + owner->nested.push_back(std::move(nested)); + } +} + +struct OwnedArray { + std::vector> columns; + std::vector children; + std::vector child_pointers; + std::vector> buffers; + std::vector> nested; + std::shared_ptr keep_alive; // what the borrowed columns point into +}; + +void ReleaseChildArray(ArrowArray* array) { array->release = nullptr; } + +void ReleaseArray(ArrowArray* array) { + delete static_cast(array->private_data); + array->release = nullptr; +} + +// Points `out` at a borrowed column's buffers, and at its dictionary where it +// has one. +void FillBorrowed(const BorrowedColumn& column, OwnedArray* owner, ArrowArray* out) { + out->length = column.length; + out->null_count = 0; + out->offset = 0; + out->n_children = 0; + out->children = nullptr; + out->dictionary = nullptr; + out->release = ReleaseChildArray; + out->private_data = nullptr; + if (column.type == ExportType::kLargeString) { + owner->buffers.push_back({nullptr, column.values, column.text}); + } else { + owner->buffers.push_back({nullptr, column.values}); + } + out->n_buffers = static_cast(owner->buffers.back().size()); + out->buffers = owner->buffers.back().data(); + if (column.type == ExportType::kDictionary) { + auto nested = std::make_unique(); + nested->children.resize(1); + nested->buffers.reserve(1); + ArrowArray* values = &nested->children[0]; + values->length = column.dictionary_size; + values->null_count = 0; + values->offset = 0; + values->n_children = 0; + values->children = nullptr; + values->dictionary = nullptr; + values->release = ReleaseChildArray; + values->private_data = nullptr; + nested->buffers.push_back( + {nullptr, column.dictionary_offsets, column.dictionary_text}); + values->n_buffers = 3; + values->buffers = nested->buffers.back().data(); + out->dictionary = values; + owner->nested.push_back(std::move(nested)); + } +} + +// Points `out` at the column's buffers, which `owner` now keeps alive. +void FillArray(BatchBuilder::Column* column, OwnedArray* owner, ArrowArray* out) { + if (column->borrowed) { + FillBorrowed(column->borrow, owner, out); + return; + } + out->length = column->length; + out->null_count = column->null_count; + out->offset = 0; + out->n_children = 0; + out->children = nullptr; + out->dictionary = nullptr; + out->release = ReleaseChildArray; + out->private_data = nullptr; + std::vector buffers; + buffers.push_back(column->validity.empty() ? nullptr : column->validity.data()); + switch (column->type) { + case ExportType::kString: + buffers.push_back(column->offsets.data()); + buffers.push_back(column->data.data()); + break; + case ExportType::kStringList: + buffers.push_back(column->offsets.data()); + break; + default: + buffers.push_back(column->data.data()); + break; + } + owner->buffers.push_back(std::move(buffers)); + out->n_buffers = static_cast(owner->buffers.back().size()); + out->buffers = owner->buffers.back().data(); + if (column->child) { + auto nested = std::make_unique(); + nested->children.resize(1); + nested->buffers.reserve(1); + FillArray(column->child.get(), nested.get(), &nested->children[0]); + nested->child_pointers = {&nested->children[0]}; + out->n_children = 1; + out->children = nested->child_pointers.data(); + owner->nested.push_back(std::move(nested)); + } +} + +} // namespace + +BatchBuilder::BatchBuilder() = default; +BatchBuilder::~BatchBuilder() = default; +BatchBuilder::BatchBuilder(BatchBuilder&&) noexcept = default; +BatchBuilder& BatchBuilder::operator=(BatchBuilder&&) noexcept = default; + +int BatchBuilder::AddColumn(std::string name, ExportType type) { + Column column; + column.name = std::move(name); + column.type = type; + if (type == ExportType::kString || type == ExportType::kStringList) { + column.offsets.push_back(0); + } + if (type == ExportType::kStringList) { + column.child = std::make_unique(); + column.child->name = "item"; + column.child->type = ExportType::kString; + column.child->offsets.push_back(0); + } + columns_.push_back(std::move(column)); + return static_cast(columns_.size() - 1); +} + +int BatchBuilder::AddBorrowed(std::string name, BorrowedColumn column) { + Column owned; + owned.name = std::move(name); + owned.type = column.type; + owned.length = column.length; + owned.borrowed = true; + owned.borrow = column; + columns_.push_back(std::move(owned)); + return static_cast(columns_.size() - 1); +} + +void BatchBuilder::KeepAlive(std::shared_ptr owner) { owner_ = std::move(owner); } + +size_t BatchBuilder::NumColumns() const { return columns_.size(); } + +std::vector BatchBuilder::ColumnNames() const { + std::vector names; + for (const Column& column : columns_) names.push_back(column.name); + return names; +} + +void BatchBuilder::AppendString(int at, std::string_view value) { + Column& column = columns_[at]; + column.data.insert(column.data.end(), value.begin(), value.end()); + column.offsets.push_back(static_cast(column.data.size())); + column.SetValid(true); +} + +void BatchBuilder::AppendBoolean(int at, bool value) { + columns_[at].PutBit(value); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendUInt8(int at, uint8_t value) { + columns_[at].PutFixed(value); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendUInt32(int at, uint32_t value) { + columns_[at].PutFixed(value); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendUInt64(int at, uint64_t value) { + columns_[at].PutFixed(value); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendDouble(int at, double value) { + columns_[at].PutFixed(value); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendDate32(int at, int32_t days) { + columns_[at].PutFixed(days); + columns_[at].SetValid(true); +} + +void BatchBuilder::AppendStringList(int at, const std::vector& values) { + Column& column = columns_[at]; + Column& items = *column.child; + for (const std::string& value : values) { + items.data.insert(items.data.end(), value.begin(), value.end()); + items.offsets.push_back(static_cast(items.data.size())); + items.SetValid(true); + } + column.offsets.push_back(static_cast(items.length)); + column.SetValid(true); +} + +void BatchBuilder::AppendNull(int at) { + Column& column = columns_[at]; + switch (column.type) { + case ExportType::kString: + column.offsets.push_back(static_cast(column.data.size())); + break; + case ExportType::kStringList: + column.offsets.push_back(static_cast(column.child->length)); + break; + case ExportType::kBoolean: + column.PutBit(false); + break; + case ExportType::kUInt8: + column.PutFixed(uint8_t{0}); + break; + case ExportType::kUInt32: + column.PutFixed(uint32_t{0}); + break; + case ExportType::kUInt64: + column.PutFixed(uint64_t{0}); + break; + case ExportType::kDouble: + column.PutFixed(0.0); + break; + case ExportType::kDate32: + column.PutFixed(int32_t{0}); + break; + case ExportType::kDictionary: + case ExportType::kLargeString: + break; // borrowed only: they take no appends + } + column.SetValid(false); +} + +int64_t BatchBuilder::Rows() const { return columns_.empty() ? 0 : columns_[0].length; } + +void BatchBuilder::ExportSchema(ArrowSchema* out) const { + auto owner = std::make_unique(); + owner->children.resize(columns_.size()); + for (size_t c = 0; c < columns_.size(); ++c) { + FillSchema(columns_[c], owner.get(), &owner->children[c]); + owner->child_pointers.push_back(&owner->children[c]); + } + out->format = "+s"; + out->name = ""; + out->metadata = nullptr; + out->flags = 0; + out->n_children = static_cast(columns_.size()); + out->children = owner->child_pointers.data(); + out->dictionary = nullptr; + out->release = ReleaseSchema; + out->private_data = owner.release(); +} + +bool BatchBuilder::ExportBatch(ArrowArray* out, std::string* error) { + const int64_t rows = Rows(); + for (const Column& column : columns_) { + if (column.length != rows) { + *error = "column \"" + column.name + "\" has " + + std::to_string(column.length) + " rows where \"" + columns_[0].name + + "\" has " + std::to_string(rows); + return false; + } + } + auto owner = std::make_unique(); + owner->children.resize(columns_.size()); + owner->buffers.reserve(columns_.size() + 1); + owner->keep_alive = owner_; + for (size_t c = 0; c < columns_.size(); ++c) { + if (columns_[c].borrowed) { + // Borrowed columns stay where they are and are exported every time. + FillBorrowed(columns_[c].borrow, owner.get(), &owner->children[c]); + owner->child_pointers.push_back(&owner->children[c]); + continue; + } + // The rows move out; the column keeps its name, type and an empty start. + auto moved = std::make_unique(); + moved->name = columns_[c].name; + moved->type = columns_[c].type; + moved->length = columns_[c].length; + moved->null_count = columns_[c].null_count; + moved->validity = std::move(columns_[c].validity); + moved->offsets = std::move(columns_[c].offsets); + moved->data = std::move(columns_[c].data); + moved->child = std::move(columns_[c].child); + columns_[c].Clear(); + columns_[c].validity.clear(); + if (columns_[c].type == ExportType::kString || + columns_[c].type == ExportType::kStringList) { + columns_[c].offsets.push_back(0); + } + if (columns_[c].type == ExportType::kStringList) { + columns_[c].child = std::make_unique(); + columns_[c].child->name = "item"; + columns_[c].child->type = ExportType::kString; + columns_[c].child->offsets.push_back(0); + } + FillArray(moved.get(), owner.get(), &owner->children[c]); + owner->child_pointers.push_back(&owner->children[c]); + owner->columns.push_back(std::move(moved)); + } + out->length = rows; + out->null_count = 0; + out->offset = 0; + out->n_buffers = 1; + owner->buffers.push_back({nullptr}); + out->buffers = owner->buffers.back().data(); + out->n_children = static_cast(columns_.size()); + out->children = owner->child_pointers.data(); + out->dictionary = nullptr; + out->release = ReleaseArray; + out->private_data = owner.release(); + return true; +} + +} // namespace cpplink diff --git a/src/cpplink/arrow_export.hpp b/src/cpplink/arrow_export.hpp new file mode 100644 index 0000000..225a7ad --- /dev/null +++ b/src/cpplink/arrow_export.hpp @@ -0,0 +1,106 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include + +#include "cpplink/arrow_c.hpp" + +namespace cpplink { + +// A record batch built column by column and handed out as the C Data structs, +// which own the memory until they are released. It is the producer side of the +// seam `batch_loader` is the consumer side of: every table cpplink writes -- the +// merged predictions, the clusters, the waterfalls, the sample -- is built here +// and handed either to the parquet writer or, in Python, to pandas, without a +// copy in between and without the builder knowing which. + +enum class ExportType : uint8_t { + kString, // utf8, 32-bit offsets + kBoolean, // bit-packed + kUInt8, + kUInt32, + kUInt64, + kDouble, + kDate32, // days since the epoch + kStringList, // list + kDictionary, // uint32 indices into a large_utf8 dictionary; borrowed only + kLargeString, // large_utf8, 64-bit offsets; borrowed only +}; + +// A column the caller already holds, exported as it sits: the buffers are +// pointed at, not copied, and the batch keeps whatever `BatchBuilder::KeepAlive` +// was given until the export is released. It is how a run's predictions leave +// as a table without a row being copied, and how the record ids come out as a +// dictionary over the store's own id arena with the rows as indices. +struct BorrowedColumn { + ExportType type = ExportType::kUInt32; + int64_t length = 0; + // Fixed-width values, the dictionary indices, or a large string's offsets. + const void* values = nullptr; + const char* text = nullptr; // a large string's bytes + // The dictionary, for kDictionary: 64-bit offsets over the text, as an id + // arena holds them. + const int64_t* dictionary_offsets = nullptr; + const char* dictionary_text = nullptr; + int64_t dictionary_size = 0; +}; + +class BatchBuilder { + public: + BatchBuilder(); + ~BatchBuilder(); + BatchBuilder(BatchBuilder&&) noexcept; + BatchBuilder& operator=(BatchBuilder&&) noexcept; + + // Declares the next column; the index is what the appenders take. Columns + // must be declared before any row is appended. + int AddColumn(std::string name, ExportType type); + // Declares a column that is already built, and stays where it is. A + // borrowed column takes no appends, is exported every time rather than moved + // out, and must be as long as the rows of the batch. + int AddBorrowed(std::string name, BorrowedColumn column); + // What every export of this builder keeps alive until it is released: the + // owner of whatever the borrowed columns point into. + void KeepAlive(std::shared_ptr owner); + size_t NumColumns() const; + std::vector ColumnNames() const; + + // One value of one column. Every column of a batch must end at the same + // length, which `ExportBatch` checks; `AppendNull` is valid on every type. + void AppendString(int column, std::string_view value); + void AppendBoolean(int column, bool value); + void AppendUInt8(int column, uint8_t value); + void AppendUInt32(int column, uint32_t value); + void AppendUInt64(int column, uint64_t value); + void AppendDouble(int column, double value); + void AppendDate32(int column, int32_t days); + void AppendStringList(int column, const std::vector& values); + void AppendNull(int column); + + // Rows appended to the first column since the last export, or a borrowed + // first column's length. + int64_t Rows() const; + + // The schema of the columns declared, as a struct with one child per column. + // The caller releases it. + void ExportSchema(ArrowSchema* out) const; + + // Moves every row appended so far out as one batch, leaving the builder empty + // and ready for the next. The caller releases the array, which frees the + // rows. Fails, leaving the rows in place, where the columns differ in length. + bool ExportBatch(ArrowArray* out, std::string* error); + + struct Column; // the buffers of one column, in the layout the C structs want + + private: + std::vector columns_; + std::shared_ptr owner_; +}; + +} // namespace cpplink diff --git a/src/cpplink/batch_loader.cpp b/src/cpplink/batch_loader.cpp new file mode 100644 index 0000000..d4e1b3a --- /dev/null +++ b/src/cpplink/batch_loader.cpp @@ -0,0 +1,986 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#include "cpplink/batch_loader.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cpplink { +namespace { + +// What a format string says, reduced to the cases the store can read. +enum class Kind { + kNull, + kBool, + kInt8, + kUInt8, + kInt16, + kUInt16, + kInt32, + kUInt32, + kInt64, + kUInt64, + kFloat, + kDouble, + kUtf8, + kLargeUtf8, + kUtf8View, + kList, + kLargeList, + kDate32, + kDate64, + kTimestamp, + kOther, +}; + +struct Format { + Kind kind = Kind::kOther; + int64_t per_day = 0; // timestamps: units in one day + char unit = 0; // timestamps: s, m, u or n + std::string_view timezone; +}; + +Format ParseFormat(const char* text) { + const std::string_view f(text == nullptr ? "" : text); + Format format; + if (f == "n") { + format.kind = Kind::kNull; + } else if (f == "b") { + format.kind = Kind::kBool; + } else if (f == "c") { + format.kind = Kind::kInt8; + } else if (f == "C") { + format.kind = Kind::kUInt8; + } else if (f == "s") { + format.kind = Kind::kInt16; + } else if (f == "S") { + format.kind = Kind::kUInt16; + } else if (f == "i") { + format.kind = Kind::kInt32; + } else if (f == "I") { + format.kind = Kind::kUInt32; + } else if (f == "l") { + format.kind = Kind::kInt64; + } else if (f == "L") { + format.kind = Kind::kUInt64; + } else if (f == "f") { + format.kind = Kind::kFloat; + } else if (f == "g") { + format.kind = Kind::kDouble; + } else if (f == "u") { + format.kind = Kind::kUtf8; + } else if (f == "U") { + format.kind = Kind::kLargeUtf8; + } else if (f == "vu") { + format.kind = Kind::kUtf8View; + } else if (f == "+l") { + format.kind = Kind::kList; + } else if (f == "+L") { + format.kind = Kind::kLargeList; + } else if (f == "tdD") { + format.kind = Kind::kDate32; + } else if (f == "tdm") { + format.kind = Kind::kDate64; + } else if (f.size() >= 4 && f[0] == 't' && f[1] == 's' && f[3] == ':') { + format.unit = f[2]; + switch (f[2]) { + case 's': + format.per_day = 86400LL; + break; + case 'm': + format.per_day = 86400000LL; + break; + case 'u': + format.per_day = 86400000000LL; + break; + case 'n': + format.per_day = 86400000000000LL; + break; + default: + return format; + } + format.kind = Kind::kTimestamp; + format.timezone = f.substr(4); + } + return format; +} + +bool IsInteger(Kind kind) { return kind >= Kind::kInt8 && kind <= Kind::kUInt64; } + +bool BitAt(const uint8_t* bits, int64_t i) { return (bits[i >> 3] >> (i & 7)) & 1; } + +// One array read at its slice: `i` below is always relative to `array->offset`. +struct ArrayView { + const ArrowSchema* schema = nullptr; + const ArrowArray* array = nullptr; + Format format; + + ArrayView(const ArrowSchema* s, const ArrowArray* a) + : schema(s), array(a), format(ParseFormat(s->format)) {} + + int64_t length() const { return array->length; } + template + const T* Buffer(int64_t index) const { + return static_cast(array->buffers[index]); + } + bool IsNull(int64_t i) const { + if (array->null_count == 0 || array->n_buffers < 1) return false; + const uint8_t* validity = Buffer(0); + return validity != nullptr && !BitAt(validity, array->offset + i); + } + // The child of a list or struct, read at this array's slice where the + // format says the parent's offset applies to it (a struct's does; a list's + // children are addressed through the offsets instead). + ArrayView Child(int64_t index) const { + return ArrayView(schema->children[index], array->children[index]); + } +}; + +// A shallow copy of an array re-sliced to `[begin, begin + length)` of its own +// values. The copy owns nothing and is never released; it exists so a struct's +// children and a list's elements can be read through the same view. +struct Slice { + ArrowArray array; + ArrayView view; + + Slice(const ArrayView& of, int64_t begin, int64_t length) + : array(*of.array), view(of.schema, &array) { + array.offset = of.array->offset + begin; + array.length = length; + } +}; + +std::string KindName(const Format& format) { + switch (format.kind) { + case Kind::kNull: + return "null"; + case Kind::kBool: + return "bool"; + case Kind::kInt8: + return "int8"; + case Kind::kUInt8: + return "uint8"; + case Kind::kInt16: + return "int16"; + case Kind::kUInt16: + return "uint16"; + case Kind::kInt32: + return "int32"; + case Kind::kUInt32: + return "uint32"; + case Kind::kInt64: + return "int64"; + case Kind::kUInt64: + return "uint64"; + case Kind::kFloat: + return "float"; + case Kind::kDouble: + return "double"; + case Kind::kUtf8: + return "string"; + case Kind::kLargeUtf8: + return "large_string"; + case Kind::kUtf8View: + return "string_view"; + case Kind::kList: + return "list"; + case Kind::kLargeList: + return "large_list"; + case Kind::kDate32: + return "date32[day]"; + case Kind::kDate64: + return "date64[ms]"; + case Kind::kTimestamp: { + std::string name = "timestamp["; + name += format.unit == 's' ? "s" + : format.unit == 'm' ? "ms" + : format.unit == 'u' ? "us" + : "ns"; + if (!format.timezone.empty()) { + name += ", tz="; + name += format.timezone; + } + return name + "]"; + } + case Kind::kOther: + break; + } + return ""; +} + +// Calls `fn(i, index)` for every non-null row of an integer array, whatever its +// width or sign, or returns false where the array is not an integer at all. +template +bool WithIntegers(const ArrayView& view, Fn&& fn) { + auto walk = [&](auto tag) { + using T = decltype(tag); + const T* values = view.Buffer(1); + for (int64_t i = 0; i < view.length(); ++i) { + if (!view.IsNull(i)) fn(i, values[view.array->offset + i]); + } + }; + switch (view.format.kind) { + case Kind::kInt8: + walk(int8_t{}); + return true; + case Kind::kUInt8: + walk(uint8_t{}); + return true; + case Kind::kInt16: + walk(int16_t{}); + return true; + case Kind::kUInt16: + walk(uint16_t{}); + return true; + case Kind::kInt32: + walk(int32_t{}); + return true; + case Kind::kUInt32: + walk(uint32_t{}); + return true; + case Kind::kInt64: + walk(int64_t{}); + return true; + case Kind::kUInt64: + walk(uint64_t{}); + return true; + default: + return false; + } +} + +// `fn(value)` for row `i` of an integer array, whatever its width or sign. +template +void WithIntegerAt(const ArrayView& view, int64_t i, Fn&& fn) { + const int64_t at = view.array->offset + i; + switch (view.format.kind) { + case Kind::kInt8: + fn(view.Buffer(1)[at]); + break; + case Kind::kUInt8: + fn(view.Buffer(1)[at]); + break; + case Kind::kInt16: + fn(view.Buffer(1)[at]); + break; + case Kind::kUInt16: + fn(view.Buffer(1)[at]); + break; + case Kind::kInt32: + fn(view.Buffer(1)[at]); + break; + case Kind::kUInt32: + fn(view.Buffer(1)[at]); + break; + case Kind::kInt64: + fn(view.Buffer(1)[at]); + break; + case Kind::kUInt64: + fn(view.Buffer(1)[at]); + break; + default: + break; + } +} + +// The bytes of value `i` of a variable-length string array, by its offsets. +template +std::string_view StringAt(const ArrayView& view, int64_t i) { + const Offset* offsets = view.Buffer(1) + view.array->offset; + const char* data = view.Buffer(2); + return std::string_view(data + offsets[i], + static_cast(offsets[i + 1] - offsets[i])); +} + +// A string view is 16 bytes: the length, then either the bytes inline or a +// prefix, the index of a data buffer and an offset into it. +std::string_view StringViewAt(const ArrayView& view, int64_t i) { + const uint8_t* entry = view.Buffer(1) + 16 * (view.array->offset + i); + int32_t length = 0; + std::memcpy(&length, entry, sizeof(length)); + if (length <= 12) + return std::string_view(reinterpret_cast(entry + 4), length); + int32_t buffer = 0; + int32_t offset = 0; + std::memcpy(&buffer, entry + 8, sizeof(buffer)); + std::memcpy(&offset, entry + 12, sizeof(offset)); + return std::string_view(view.Buffer(2 + buffer) + offset, length); +} + +template +bool ForEachText(const ArrayView& view, Visit&& visit, std::string* error); + +// The dictionary's values as owned text, so an integer dictionary spells its +// digits once rather than per row. +bool DictionaryValues(const ArrayView& view, std::vector* values, + std::vector* valid, std::string* error) { + const ArrayView dictionary(view.schema->dictionary, view.array->dictionary); + return ForEachText( + dictionary, + [&](std::string_view value, bool is_null) { + values->emplace_back(value); + valid->push_back(!is_null); + }, + error); +} + +// Visits every value of a text-like column in order as `visit(value, is_null)`. +// A frame rarely holds the types the schema names: pandas and polars hand over +// `large_string` or `string_view` for text, and a postcode, a year or an id that +// was numeric arrives as an integer. Both are the string the user meant, so an +// integer is visited as its decimal digits and interned like any other value, +// which is what makes "2000" in one input and 2000 in another the same id. +template +bool ForEachText(const ArrayView& view, Visit&& visit, std::string* error) { + if (view.schema->dictionary != nullptr) { + std::vector values; + std::vector valid; + if (!DictionaryValues(view, &values, &valid, error)) return false; + int64_t next = 0; + const bool integer = WithIntegers(view, [&](int64_t i, auto index) { + for (; next < i; ++next) visit(std::string_view(), true); + next = i + 1; + const auto at = static_cast(index); + visit(std::string_view(values[at]), !valid[at]); + }); + if (!integer) { + *error = "dictionary indices are not integers"; + return false; + } + for (; next < view.length(); ++next) visit(std::string_view(), true); + return true; + } + switch (view.format.kind) { + case Kind::kUtf8: + case Kind::kLargeUtf8: + case Kind::kUtf8View: + for (int64_t i = 0; i < view.length(); ++i) { + if (view.IsNull(i)) { + visit(std::string_view(), true); + } else if (view.format.kind == Kind::kUtf8) { + visit(StringAt(view, i), false); + } else if (view.format.kind == Kind::kLargeUtf8) { + visit(StringAt(view, i), false); + } else { + visit(StringViewAt(view, i), false); + } + } + return true; + case Kind::kNull: + for (int64_t i = 0; i < view.length(); ++i) visit(std::string_view(), true); + return true; + default: + break; + } + int64_t next = 0; + char digits[24]; // -9223372036854775808 is 20 characters + const bool integer = WithIntegers(view, [&](int64_t i, auto value) { + for (; next < i; ++next) visit(std::string_view(), true); + next = i + 1; + const auto result = std::to_chars(digits, digits + sizeof(digits), value); + visit(std::string_view(digits, result.ptr - digits), false); + }); + if (integer) { + for (; next < view.length(); ++next) visit(std::string_view(), true); + return true; + } + *error = "expected a string or integer column, found " + ArrowTypeName(*view.schema); + return false; +} + +bool AppendStrings(const ArrayView& view, StringColumn* column, std::string* error) { + if (view.schema->dictionary != nullptr) { + // A category column is the store's own shape: each distinct value is + // interned once, on the first row that uses it, and the rows are a remap + // of the indices. A category no row holds is not interned, so the + // dictionary is what the rows say, as it is when the same column is read + // from a file. + std::vector values; + std::vector valid; + if (!DictionaryValues(view, &values, &valid, error)) return false; + std::vector mapped(values.size(), kNullId); + std::vector seen(values.size(), false); + const size_t first = column->ids.size(); + column->ids.resize(first + static_cast(view.length()), kNullId); + const bool integer = WithIntegers(view, [&](int64_t i, auto index) { + const auto at = static_cast(index); + if (!seen[at]) { + seen[at] = true; + if (valid[at]) mapped[at] = column->dict.Intern(values[at]); + } + column->ids[first + i] = mapped[at]; + }); + if (!integer) { + *error = "dictionary indices are not integers"; + return false; + } + return true; + } + return ForEachText( + view, + [&](std::string_view value, bool is_null) { + column->ids.push_back(is_null ? kNullId : column->dict.Intern(value)); + }, + error); +} + +// Slices the interned elements of a list array into one sorted, deduplicated cell +// per row, whether the offsets are 32-bit (`list`) or 64-bit (`large_list`). +template +void AppendListRows(const ArrayView& view, const std::vector& elements, + Offset first, StringListColumn* column) { + const Offset* offsets = view.Buffer(1) + view.array->offset; + std::vector row; + for (int64_t i = 0; i < view.length(); ++i) { + row.clear(); + if (!view.IsNull(i)) { + for (Offset j = offsets[i] - first; j < offsets[i + 1] - first; ++j) { + if (elements[j] != kNullId) row.push_back(elements[j]); + } + } + // Sorted and deduplicated here so overlap is a linear merge later. + std::sort(row.begin(), row.end()); + row.erase(std::unique(row.begin(), row.end()), row.end()); + column->ids.insert(column->ids.end(), row.begin(), row.end()); + column->offsets.push_back(column->ids.size()); + } +} + +template +bool AppendLists(const ArrayView& view, StringListColumn* column, std::string* error) { + // The elements are interned first, as one flat run over the child array, so a + // list of integers goes through the same visitor a scalar column does; the rows + // are then cut out of that run by the list's own offsets. Only the elements + // this slice reaches are interned, or a sliced frame would leave values in the + // dictionary that no row holds. + const Offset* offsets = view.Buffer(1) + view.array->offset; + const Offset first = view.length() > 0 ? offsets[0] : 0; + const Offset last = view.length() > 0 ? offsets[view.length()] : 0; + const Slice elements_slice(view.Child(0), first, last - first); + std::vector elements; + elements.reserve(static_cast(last - first)); + std::string element_error; + const bool ok = ForEachText( + elements_slice.view, + [&](std::string_view value, bool is_null) { + elements.push_back(is_null ? kNullId : column->dict.Intern(value)); + }, + &element_error); + if (!ok) { + *error = "expected a list of strings or integers, found a list of " + + ArrowTypeName(*view.schema->children[0]); + return false; + } + if (column->offsets.empty()) column->offsets.push_back(0); + AppendListRows(view, elements, first, column); + return true; +} + +bool AppendStringLists(const ArrayView& view, StringListColumn* column, + std::string* error) { + switch (view.format.kind) { + case Kind::kList: + return AppendLists(view, column, error); + case Kind::kLargeList: + return AppendLists(view, column, error); + case Kind::kNull: + if (column->offsets.empty()) column->offsets.push_back(0); + for (int64_t i = 0; i < view.length(); ++i) { + column->offsets.push_back(column->ids.size()); + } + return true; + default: + *error = "expected a list column, found " + ArrowTypeName(*view.schema); + return false; + } +} + +// Days since the Unix epoch, whatever the column's temporal unit. +bool AppendDates(const ArrayView& view, DateColumn* column, std::string* error) { + const int64_t offset = view.array->offset; + switch (view.format.kind) { + case Kind::kDate32: { + const int32_t* values = view.Buffer(1) + offset; + for (int64_t i = 0; i < view.length(); ++i) { + column->values.push_back(view.IsNull(i) ? kNullDate : values[i]); + } + return true; + } + case Kind::kDate64: + case Kind::kTimestamp: { + const int64_t per_day = + view.format.kind == Kind::kDate64 ? 86400000LL : view.format.per_day; + const int64_t* values = view.Buffer(1) + offset; + for (int64_t i = 0; i < view.length(); ++i) { + column->values.push_back(view.IsNull(i) + ? kNullDate + : static_cast(values[i] / per_day)); + } + return true; + } + case Kind::kNull: + column->values.insert(column->values.end(), view.length(), kNullDate); + return true; + default: + *error = "expected a date or timestamp column, found " + + ArrowTypeName(*view.schema); + return false; + } +} + +bool AppendDoubles(const ArrayView& view, DoubleColumn* column, std::string* error) { + const double kMissing = std::numeric_limits::quiet_NaN(); + const int64_t offset = view.array->offset; + switch (view.format.kind) { + case Kind::kDouble: { + const double* values = view.Buffer(1) + offset; + for (int64_t i = 0; i < view.length(); ++i) { + column->values.push_back(view.IsNull(i) ? kMissing : values[i]); + } + return true; + } + case Kind::kFloat: { + const float* values = view.Buffer(1) + offset; + for (int64_t i = 0; i < view.length(); ++i) { + column->values.push_back(view.IsNull(i) ? kMissing : values[i]); + } + return true; + } + case Kind::kNull: + column->values.insert(column->values.end(), view.length(), kMissing); + return true; + default: + *error = + "expected a floating point column, found " + ArrowTypeName(*view.schema); + return false; + } +} + +bool AppendBooleans(const ArrayView& view, BooleanColumn* column, std::string* error) { + switch (view.format.kind) { + case Kind::kBool: { + const uint8_t* bits = view.Buffer(1); + for (int64_t i = 0; i < view.length(); ++i) { + column->values.push_back(view.IsNull(i) ? kNullBoolean + : BitAt(bits, view.array->offset + i) + ? int8_t{1} + : int8_t{0}); + } + return true; + } + case Kind::kNull: + column->values.insert(column->values.end(), view.length(), kNullBoolean); + return true; + default: + *error = "expected a boolean column, found " + ArrowTypeName(*view.schema); + return false; + } +} + +bool AppendIds(const ArrayView& view, IdColumn* column, std::string* error) { + std::string text_error; + const bool ok = ForEachText( + view, [&](std::string_view value, bool) { column->Append(value); }, &text_error); + if (ok) return true; + *error = "expected a string or integer unique_id column, found " + + ArrowTypeName(*view.schema); + return false; +} + +bool AppendColumn(const ArrayView& view, ColumnType type, Column* column, + std::string* error) { + switch (type) { + case ColumnType::kString: + return AppendStrings(view, &std::get(*column), error); + case ColumnType::kStringList: + return AppendStringLists(view, &std::get(*column), error); + case ColumnType::kDate: + return AppendDates(view, &std::get(*column), error); + case ColumnType::kBoolean: + return AppendBooleans(view, &std::get(*column), error); + case ColumnType::kDouble: + return AppendDoubles(view, &std::get(*column), error); + } + return false; +} + +// The types the store does not read, named as Arrow prints them where the +// spelling is short, and by their format string otherwise. +std::string OtherTypeName(const ArrowSchema& field) { + const std::string_view f(field.format == nullptr ? "" : field.format); + if (f == "z") return "binary"; + if (f == "Z") return "large_binary"; + if (f == "vz") return "binary_view"; + if (f == "e") return "halffloat"; + if (f == "+s") return "struct"; + if (f == "+m") return "map"; + if (f.size() > 2 && f[0] == 'd' && f[1] == ':') + return "decimal(" + std::string(f.substr(2)) + ")"; + if (f.size() > 2 && f[0] == 'w' && f[1] == ':') + return "fixed_size_binary[" + std::string(f.substr(2)) + "]"; + if (f == "tts") return "time32[s]"; + if (f == "ttm") return "time32[ms]"; + if (f == "ttu") return "time64[us]"; + if (f == "ttn") return "time64[ns]"; + return std::string(f); +} + +} // namespace + +int FieldIndex(const ArrowSchema& schema, std::string_view name) { + for (int64_t i = 0; i < schema.n_children; ++i) { + const char* child = schema.children[i]->name; + if (child != nullptr && name == child) return static_cast(i); + } + return -1; +} + +bool ColumnTypeOf(const ArrowSchema& field, ColumnType* out) { + if (field.dictionary != nullptr) return ColumnTypeOf(*field.dictionary, out); + const Format format = ParseFormat(field.format); + if (IsInteger(format.kind)) { + *out = ColumnType::kString; + return true; + } + switch (format.kind) { + case Kind::kUtf8: + case Kind::kLargeUtf8: + case Kind::kUtf8View: + *out = ColumnType::kString; + return true; + case Kind::kList: + case Kind::kLargeList: + *out = ColumnType::kStringList; + return true; + case Kind::kDate32: + case Kind::kDate64: + case Kind::kTimestamp: + *out = ColumnType::kDate; + return true; + case Kind::kBool: + *out = ColumnType::kBoolean; + return true; + case Kind::kFloat: + case Kind::kDouble: + *out = ColumnType::kDouble; + return true; + default: + return false; + } +} + +std::string ArrowTypeName(const ArrowSchema& field) { + const Format format = ParseFormat(field.format); + if (field.dictionary != nullptr) { + return "dictionary"; + } + if (format.kind == Kind::kList || format.kind == Kind::kLargeList) { + const ArrowSchema& child = *field.children[0]; + const std::string name = child.name == nullptr ? "item" : child.name; + return KindName(format) + "<" + name + ": " + ArrowTypeName(child) + ">"; + } + if (format.kind == Kind::kOther) return OtherTypeName(field); + return KindName(format); +} + +bool ResolveBatchColumns(const ArrowSchema& schema, const Schema& spec, + BatchColumns* columns, std::string* error) { + columns->id = -1; + columns->children.assign(spec.columns.size(), -1); + if (!spec.unique_id.empty()) { + columns->id = FieldIndex(schema, spec.unique_id); + if (columns->id < 0) { + *error = "unique_id column \"" + spec.unique_id + "\" is not in the input"; + return false; + } + } + // A derived column is not among them: it is computed from another column when + // the store is finalized, and the input is not expected to hold it. + for (size_t i = 0; i < spec.columns.size(); ++i) { + if (spec.columns[i].IsDerived()) continue; + columns->children[i] = FieldIndex(schema, spec.columns[i].name); + if (columns->children[i] < 0) { + *error = "column \"" + spec.columns[i].name + "\" is not in the input"; + return false; + } + } + return true; +} + +bool ResolveColumnTypesFrom(const ArrowSchema& schema, const std::string& source, + Schema* spec, std::string* error) { + for (ColumnSpec& column : spec->columns) { + if (column.type_declared || column.IsDerived()) continue; + const int index = FieldIndex(schema, column.name); + if (index < 0) { + *error = "column \"" + column.name + "\" is not in " + source; + return false; + } + const ArrowSchema& field = *schema.children[index]; + if (!ColumnTypeOf(field, &column.type)) { + *error = "column \"" + column.name + "\" is " + ArrowTypeName(field) + + " in " + source + ", which cpplink cannot read; cast it to a string"; + return false; + } + } + return CheckTypes(spec, error); +} + +bool AppendRecordBatch(const ArrowSchema& schema, const ArrowArray& batch, + const Schema& spec, RecordStore* store, std::string* error) { + if (batch.n_children != schema.n_children) { + *error = "the batch holds " + std::to_string(batch.n_children) + + " columns where its schema names " + std::to_string(schema.n_children); + return false; + } + BatchColumns columns; + if (!ResolveBatchColumns(schema, spec, &columns, error)) return false; + + // A record batch is a struct array, and a struct's offset applies to its + // children: each is read at the batch's slice, whatever its own length says. + const ArrayView root(&schema, &batch); + auto column_view = [&](int child) { + return Slice(root.Child(child), batch.offset, batch.length); + }; + for (int64_t i = 0; i < batch.n_children; ++i) { + if (batch.children[i]->length < batch.offset + batch.length) { + *error = std::string("column \"") + schema.children[i]->name + "\" holds " + + std::to_string(batch.children[i]->length) + + " rows where the batch has " + std::to_string(batch.length); + return false; + } + } + + if (columns.id >= 0) { + const Slice ids = column_view(columns.id); + if (!AppendIds(ids.view, &store->mutable_ids(), error)) { + *error = spec.unique_id + ": " + *error; + return false; + } + } + for (size_t i = 0; i < spec.columns.size(); ++i) { + if (columns.children[i] < 0) continue; + const Slice column = column_view(columns.children[i]); + if (!AppendColumn(column.view, spec.columns[i].type, &store->mutable_column(i), + error)) { + *error = spec.columns[i].name + ": " + *error; + return false; + } + } + return true; +} + +bool AppendArrayStream(ArrowArrayStream* stream, const Schema& spec, RecordStore* store, + uint64_t* rows, int* batches, std::string* error) { + auto fail = [&](const std::string& what, int code) { + const char* detail = + stream->get_last_error == nullptr ? nullptr : stream->get_last_error(stream); + *error = + what + ": " + (detail == nullptr ? "error " + std::to_string(code) : detail); + if (stream->release != nullptr) stream->release(stream); + return false; + }; + ArrowSchema schema; + schema.release = nullptr; + if (const int code = stream->get_schema(stream, &schema); code != 0) { + return fail("cannot read the stream's schema", code); + } + bool ok = true; + while (ok) { + ArrowArray batch; + batch.release = nullptr; + if (const int code = stream->get_next(stream, &batch); code != 0) { + schema.release(&schema); + return fail("cannot read the next batch", code); + } + if (batch.release == nullptr) break; // the stream has ended + ok = AppendRecordBatch(schema, batch, spec, store, error); + if (ok) { + *rows += static_cast(batch.length); + ++*batches; + } + batch.release(&batch); + } + schema.release(&schema); + if (stream->release != nullptr) stream->release(stream); + return ok; +} + +bool LoadStreams(const std::vector& inputs, const Schema& spec, + RecordStore* store, LoadStats* stats, std::string* error) { + const auto started = std::chrono::steady_clock::now(); + if (inputs.empty()) { + *error = "no input was given"; + return false; + } + // Boundaries are recorded as they are crossed, so a dataset is the row range + // between two of them and nothing per row has to be stored. + std::vector starts; + std::vector per_input; + std::vector names; + starts.push_back(0); + uint64_t total = 0; + int batches = 0; + bool ok = true; + for (size_t i = 0; i < inputs.size(); ++i) { + uint64_t rows = 0; + int count = 0; + if (ok) { + ok = AppendArrayStream(inputs[i].stream, spec, store, &rows, &count, error); + // With several inputs the message has to say which one failed; a + // column missing from the second reads the same as one missing from + // the first otherwise. + if (!ok && inputs.size() > 1) *error = inputs[i].name + ": " + *error; + } else if (inputs[i].stream->release != nullptr) { + inputs[i].stream->release(inputs[i].stream); + } + total += rows; + batches += count; + per_input.push_back(rows); + starts.push_back(total); + names.push_back(inputs[i].name); + } + if (!ok) return false; + + store->set_num_records(total); + // One dataset is the absence of a boundary, not a boundary at each end: the + // dedup path then pays no lookup at all. + if (inputs.size() > 1) { + store->set_datasets(std::move(starts)); + store->set_dataset_names(std::move(names)); + } + store->Finalize(); + + if (stats != nullptr) { + stats->rows = total; + stats->row_groups = batches; + stats->dataset_rows = std::move(per_input); + stats->seconds = + std::chrono::duration(std::chrono::steady_clock::now() - started) + .count(); + } + return true; +} + +struct TextReader::Impl { + ArrowArray sliced; + ArrayView view; + std::vector values; // a dictionary's, decoded once + std::vector valid; + Impl(const ArrowSchema* schema, const ArrowArray* array) + : sliced(*array), view(schema, &sliced) {} +}; + +TextReader::TextReader() = default; +TextReader::~TextReader() = default; + +bool TextReader::Bind(const ArrowSchema& schema, const ArrowArray& batch, int column, + std::string* error) { + impl_ = std::make_unique(schema.children[column], batch.children[column]); + impl_->sliced.offset += batch.offset; + impl_->sliced.length = batch.length; + const Format& format = impl_->view.format; + if (impl_->view.schema->dictionary != nullptr) { + if (!IsInteger(format.kind)) { + *error = "dictionary indices are not integers"; + return false; + } + return DictionaryValues(impl_->view, &impl_->values, &impl_->valid, error); + } + if (format.kind == Kind::kUtf8 || format.kind == Kind::kLargeUtf8 || + format.kind == Kind::kUtf8View || format.kind == Kind::kNull || + IsInteger(format.kind)) { + return true; + } + *error = "expected a string or integer column, found " + ArrowTypeName(schema); + return false; +} + +bool TextReader::At(int64_t row, std::string_view* out, char (*scratch)[24]) const { + const ArrayView& view = impl_->view; + if (view.IsNull(row)) return false; + if (view.schema->dictionary != nullptr) { + int64_t index = 0; + WithIntegerAt(view, row, + [&](auto value) { index = static_cast(value); }); + if (!impl_->valid[static_cast(index)]) return false; + *out = impl_->values[static_cast(index)]; + return true; + } + switch (view.format.kind) { + case Kind::kUtf8: + *out = StringAt(view, row); + return true; + case Kind::kLargeUtf8: + *out = StringAt(view, row); + return true; + case Kind::kUtf8View: + *out = StringViewAt(view, row); + return true; + case Kind::kNull: + return false; + default: + break; + } + WithIntegerAt(view, row, [&](auto value) { + const auto result = std::to_chars(*scratch, *scratch + sizeof(*scratch), value); + *out = std::string_view(*scratch, result.ptr - *scratch); + }); + return true; +} + +struct NumberReader::Impl { + ArrowArray sliced; + ArrayView view; + Impl(const ArrowSchema* schema, const ArrowArray* array) + : sliced(*array), view(schema, &sliced) {} +}; + +NumberReader::NumberReader() = default; +NumberReader::~NumberReader() = default; + +bool NumberReader::Bind(const ArrowSchema& schema, const ArrowArray& batch, int column, + std::string* error) { + impl_ = std::make_unique(schema.children[column], batch.children[column]); + impl_->sliced.offset += batch.offset; + impl_->sliced.length = batch.length; + const Kind kind = impl_->view.format.kind; + if (kind == Kind::kDouble || kind == Kind::kFloat || kind == Kind::kNull || + IsInteger(kind)) { + return true; + } + *error = "expected a numeric column, found " + ArrowTypeName(schema); + return false; +} + +bool NumberReader::At(int64_t row, double* out) const { + const ArrayView& view = impl_->view; + if (view.IsNull(row)) return false; + const int64_t at = view.array->offset + row; + switch (view.format.kind) { + case Kind::kDouble: + *out = view.Buffer(1)[at]; + return true; + case Kind::kFloat: + *out = view.Buffer(1)[at]; + return true; + case Kind::kNull: + return false; + default: + break; + } + WithIntegerAt(view, row, [&](auto value) { *out = static_cast(value); }); + return true; +} + +} // namespace cpplink diff --git a/src/cpplink/batch_loader.hpp b/src/cpplink/batch_loader.hpp new file mode 100644 index 0000000..7cf1356 --- /dev/null +++ b/src/cpplink/batch_loader.hpp @@ -0,0 +1,130 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include +#include +#include + +#include "cpplink/arrow_c.hpp" +#include "cpplink/record_store.hpp" +#include "cpplink/schema.hpp" + +namespace cpplink { + +// Arrow record batches, through the C Data Interface, appended to the store. +// +// This is the one decoder of columns into the store. The parquet loader exports +// each row group through it, and a data frame handed over by Python arrives here +// directly, so the two front ends cannot read a value differently: a `large_string` +// or an integer column is text either way, a timestamp is days either way. Nothing +// here links Arrow C++; the structs are the specification's and the buffers are +// read where the producer left them, which is what keeps a frame from being +// copied on its way in. Every string is interned, so the store is a re-encoding +// rather than a view, and the only bytes duplicated are the dictionary's. + +struct LoadStats { + uint64_t rows = 0; + int row_groups = 0; // row groups of parquet, or batches of a stream + double seconds = 0.0; + std::vector dataset_rows; // rows contributed by each input, in order +}; + +// The column type an Arrow field reads as, or false where none does. An integer +// is text: a postcode, a year or a phone number that was numeric in the frame is +// still the value it spells, and a writer who wanted a double wrote one. A +// dictionary-encoded field reads as its values' type. +bool ColumnTypeOf(const ArrowSchema& field, ColumnType* out); + +// The field's type spelled the way Arrow C++ prints it (`int64`, `large_string`, +// `list`, `timestamp[ns, tz=UTC]`), for messages and reports. +std::string ArrowTypeName(const ArrowSchema& field); + +// Which child of a batch each store column reads, by name. `children[i]` is the +// child for `spec.columns[i]`, -1 for a derived column, which no input holds; +// `id` is the child for `spec.unique_id`, -1 where the schema names none. Fails +// naming the first column the batch lacks. +struct BatchColumns { + int id = -1; + std::vector children; +}; +bool ResolveBatchColumns(const ArrowSchema& schema, const Schema& spec, + BatchColumns* columns, std::string* error); + +// Fills in the type of every column the schema left undeclared from the field +// the Arrow schema holds, then runs the checks those types decide. `source` +// names the file or frame in the message. Must run before a store is built from +// the schema, because the store's column layout is the types. +bool ResolveColumnTypesFrom(const ArrowSchema& schema, const std::string& source, + Schema* spec, std::string* error); + +// Appends one batch's rows to the store, column by column, without finalizing it: +// more batches may be coming, and the term frequencies are only right once every +// row is in. The batch's fields are matched to the store's columns by name. A +// sliced array (a non-zero `offset`) is read at its slice, and a dictionary- +// encoded column interns each dictionary value once and maps the indices. +bool AppendRecordBatch(const ArrowSchema& schema, const ArrowArray& batch, + const Schema& spec, RecordStore* store, std::string* error); + +// Drains a stream into the store, batch by batch, releasing each before the next +// is fetched so only one batch's buffers are ever resident beside the store. The +// stream itself is released on return, success or not. `rows` and `batches` +// count what was read. +bool AppendArrayStream(ArrowArrayStream* stream, const Schema& spec, RecordStore* store, + uint64_t* rows, int* batches, std::string* error); + +// One input of a run: a stream of batches and the name its dataset gets. +struct InputStream { + std::string name; + ArrowArrayStream* stream = nullptr; +}; + +// Loads every input in order into one store and finalizes it. Each input becomes +// a dataset, and because rows are appended one input at a time a dataset is a +// contiguous row range -- which is what makes a link mode's cross-product a +// nested loop over two ranges rather than a per-row test. A single input gets +// no dataset name, because nothing needs qualifying. Every stream is released, +// success or not. +bool LoadStreams(const std::vector& inputs, const Schema& spec, + RecordStore* store, LoadStats* stats, std::string* error); + +// The child of a struct schema called `name`, or -1. +int FieldIndex(const ArrowSchema& schema, std::string_view name); + +// One text column of a batch read by row, for the files that are not the store: +// a prediction file's ids and datasets. Binds a string, string view, integer or +// dictionary-encoded column, sliced as the batch is. +class TextReader { + public: + TextReader(); + ~TextReader(); + bool Bind(const ArrowSchema& schema, const ArrowArray& batch, int column, + std::string* error); + // The value at `row`, or false where it is null. An integer's digits are + // written into `scratch`, which the view then points at. + bool At(int64_t row, std::string_view* out, char (*scratch)[24]) const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +// One numeric column of a batch read by row: a float of either width or an +// integer of any, as a double. +class NumberReader { + public: + NumberReader(); + ~NumberReader(); + bool Bind(const ArrowSchema& schema, const ArrowArray& batch, int column, + std::string* error); + bool At(int64_t row, double* out) const; + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace cpplink diff --git a/src/cpplink/cluster.cpp b/src/cpplink/cluster.cpp index a494485..ea904b9 100644 --- a/src/cpplink/cluster.cpp +++ b/src/cpplink/cluster.cpp @@ -17,14 +17,11 @@ #include #include -#include -#include -#include -#include - +#include "cpplink/arrow_export.hpp" #include "cpplink/format.hpp" #include "cpplink/id_index.hpp" #include "cpplink/merge_edges.hpp" +#include "cpplink/parquet_io.hpp" #include "cpplink/predict.hpp" namespace cpplink { @@ -214,119 +211,38 @@ bool ReadCsvEdges(const std::string& path, const RecordStore& store, const IdInd return true; } -// One row group at a time, and only the columns clustering reads -- the ids, the -// weight, and the datasets where the file has them: the pattern and the -// posterior are in the file for whoever else opens it. +// Every row of the merged parquet, through the same reader a data frame of +// predictions goes through. A row with a null id or dataset resolves to nothing +// and is counted as unresolved, as a csv row naming an unknown record is. +EdgeVisitor ResolvingVisitor(const RecordStore& store, const IdIndex& index, + EdgeTally* tally) { + return [&store, &index, tally](const EdgeRow& edge, std::string*) { + uint32_t a = 0; + uint32_t b = 0; + if (!Resolve(index, store, edge.dataset_a, edge.id_a, &a, tally) || + !Resolve(index, store, edge.dataset_b, edge.id_b, &b, tally)) { + ++tally->read; + ++tally->unresolved; + return true; + } + tally->Use(a, b, edge.weight); + return true; + }; +} + bool ReadParquetEdges(const std::string& path, const RecordStore& store, const IdIndex& index, EdgeTally* tally, std::string* error) { - auto input = arrow::io::ReadableFile::Open(path); - if (!input.ok()) { - *error = "cluster: cannot open " + path + ": " + input.status().message(); - return false; - } - parquet::arrow::FileReaderBuilder builder; - arrow::Status status = builder.Open(*input); - if (!status.ok()) { - *error = - "cluster: cannot read parquet metadata for " + path + ": " + status.message(); - return false; - } - auto reader_result = builder.Build(); - if (!reader_result.ok()) { - *error = "cluster: cannot open parquet reader for " + path + ": " + - reader_result.status().message(); - return false; - } - std::unique_ptr reader = std::move(*reader_result); - - std::shared_ptr schema; - status = reader->GetSchema(&schema); - if (!status.ok()) { - *error = "cluster: cannot read the schema of " + path + ": " + status.message(); - return false; - } - std::vector indices; - for (const char* name : {"id_a", "id_b", "match_weight"}) { - const int at = schema->GetFieldIndex(name); - if (at < 0) { - *error = "cluster: '" + path + "' has no column \"" + name + - "\", so it is not a cpplink prediction file"; - return false; - } - indices.push_back(at); - } - const bool datasets = schema->GetFieldIndex("dataset_a") >= 0; - if (datasets) { - for (const char* name : {"dataset_a", "dataset_b"}) { - const int at = schema->GetFieldIndex(name); - if (at < 0) { - *error = "cluster: '" + path + "' has dataset_a but no " + name; - return false; - } - indices.push_back(at); - } - } + EdgeColumns columns; + return ReadPredictionFile(path, "cluster", {"match_weight"}, &columns, + ResolvingVisitor(store, index, tally), error); +} - for (int group = 0; group < reader->num_row_groups(); ++group) { - auto group_result = reader->ReadRowGroup(group, indices); - if (!group_result.ok()) { - *error = "cluster: cannot read row group " + std::to_string(group) + " of " + - path + ": " + group_result.status().message(); - return false; - } - const std::shared_ptr table = *group_result; - const auto ids_a = table->GetColumnByName("id_a"); - const auto ids_b = table->GetColumnByName("id_b"); - const auto weights = table->GetColumnByName("match_weight"); - const auto sets_a = datasets ? table->GetColumnByName("dataset_a") : nullptr; - const auto sets_b = datasets ? table->GetColumnByName("dataset_b") : nullptr; - for (int chunk = 0; chunk < ids_a->num_chunks(); ++chunk) { - const auto a_array = - std::dynamic_pointer_cast(ids_a->chunk(chunk)); - const auto b_array = - std::dynamic_pointer_cast(ids_b->chunk(chunk)); - const auto weight_array = - std::dynamic_pointer_cast(weights->chunk(chunk)); - std::shared_ptr set_a; - std::shared_ptr set_b; - if (datasets) { - set_a = - std::dynamic_pointer_cast(sets_a->chunk(chunk)); - set_b = - std::dynamic_pointer_cast(sets_b->chunk(chunk)); - } - if (a_array == nullptr || b_array == nullptr || weight_array == nullptr || - (datasets && (set_a == nullptr || set_b == nullptr))) { - *error = "cluster: '" + path + - "' holds id_a, id_b, dataset_a, dataset_b or match_weight in " - "a type a cpplink prediction file does not use"; - return false; - } - for (int64_t row = 0; row < a_array->length(); ++row) { - if (a_array->IsNull(row) || b_array->IsNull(row) || - weight_array->IsNull(row) || - (datasets && (set_a->IsNull(row) || set_b->IsNull(row)))) { - ++tally->read; - ++tally->unresolved; - continue; - } - const std::string_view dataset_a = - datasets ? set_a->GetView(row) : std::string_view(); - const std::string_view dataset_b = - datasets ? set_b->GetView(row) : std::string_view(); - uint32_t a = 0; - uint32_t b = 0; - if (!Resolve(index, store, dataset_a, a_array->GetView(row), &a, tally) || - !Resolve(index, store, dataset_b, b_array->GetView(row), &b, tally)) { - ++tally->read; - ++tally->unresolved; - continue; - } - tally->Use(a, b, weight_array->Value(row)); - } - } - } - return true; +// A prediction table from wherever it came, read as the merged file is. +bool ReadStreamEdges(ArrowArrayStream* stream, const RecordStore& store, + const IdIndex& index, EdgeTally* tally, std::string* error) { + EdgeColumns columns; + return ReadPredictionStream(stream, "cluster", "predictions", {"match_weight"}, + &columns, ResolvingVisitor(store, index, tally), error); } } // namespace @@ -368,7 +284,25 @@ bool Cluster(const RecordStore& store, const ClusterOptions& options, std::string edge_file; double index_seconds = 0.0; std::error_code ec; - if (std::filesystem::is_directory(options.edge_path, ec)) { + if (options.edges != nullptr) { + // The run's own rows: nothing to resolve. + for (uint64_t i = 0; i < options.edges->Size(); ++i) { + tally.Use(options.edges->a[i], options.edges->b[i], options.edges->weight[i]); + } + } else if (options.stream != nullptr) { + const auto index_started = std::chrono::steady_clock::now(); + const IdIndex index(store); + index_seconds = std::chrono::duration(std::chrono::steady_clock::now() - + index_started) + .count(); + if (!index.Unique(error)) { + *error = "cluster: " + *error; + if (options.stream->release != nullptr) + options.stream->release(options.stream); + return false; + } + if (!ReadStreamEdges(options.stream, store, index, &tally, error)) return false; + } else if (std::filesystem::is_directory(options.edge_path, ec)) { if (!CollectShards(options.edge_path, &shards, error)) return false; for (const std::string& path : shards) { if (!ReadBinaryShard(path, &tally, error)) return false; @@ -529,84 +463,51 @@ bool WriteClustersCsv(const ClusterAssignment& assignment, const RecordStore& st } // The same columns as the csv, typed: the names as strings and the size as a -// uint64. +// uint64, built as C Data batches of a million rows and handed to the writer. bool WriteClustersParquet(const ClusterAssignment& assignment, const RecordStore& store, const ClusterOptions& options, uint64_t* written, std::string* error) { constexpr int64_t kRowGroup = 1 << 20; const bool datasets = store.NumDatasets() > 1; - std::vector> fields; - if (datasets) fields.push_back(arrow::field("dataset", arrow::utf8())); - fields.push_back(arrow::field("unique_id", arrow::utf8())); - fields.push_back(arrow::field("cluster_id", arrow::utf8())); - fields.push_back(arrow::field("cluster_size", arrow::uint64())); - const auto schema = arrow::schema(fields); - auto sink = arrow::io::FileOutputStream::Open(options.out_path); - if (!sink.ok()) { - *error = "cluster: cannot create '" + options.out_path + - "': " + sink.status().message(); + BatchBuilder builder; + const int dataset = datasets ? builder.AddColumn("dataset", ExportType::kString) : -1; + const int unique_id = builder.AddColumn("unique_id", ExportType::kString); + const int cluster_id = builder.AddColumn("cluster_id", ExportType::kString); + const int cluster_size = builder.AddColumn("cluster_size", ExportType::kUInt64); + + ParquetWriter writer; + ArrowSchema schema; + builder.ExportSchema(&schema); + const bool opened = writer.Open(options.out_path, schema, error); + schema.release(&schema); + if (!opened) { + *error = "cluster: " + *error; return false; } - auto props = parquet::WriterProperties::Builder() - .compression(parquet::Compression::SNAPPY) - ->build(); - auto opened = parquet::arrow::FileWriter::Open(*schema, arrow::default_memory_pool(), - *sink, props); - if (!opened.ok()) { - *error = "cluster: cannot open a parquet writer for '" + options.out_path + - "': " + opened.status().message(); - return false; - } - std::unique_ptr writer = std::move(*opened); - - arrow::StringBuilder dataset, unique_id, cluster_id; - arrow::UInt64Builder cluster_size; - int64_t pending = 0; auto flush = [&]() -> bool { - if (pending == 0) return true; - std::vector> arrays; - arrow::Status status; - auto take = [&](arrow::ArrayBuilder* builder) { - std::shared_ptr array; - status &= builder->Finish(&array); - arrays.push_back(std::move(array)); - }; - if (datasets) take(&dataset); - take(&unique_id); - take(&cluster_id); - take(&cluster_size); - if (status.ok()) { - status = - writer->WriteTable(*arrow::Table::Make(schema, arrays, pending), pending); - } - if (!status.ok()) { - *error = "cluster: writing a row group to '" + options.out_path + - "': " + status.message(); + if (builder.Rows() == 0) return true; + ArrowArray batch; + if (!builder.ExportBatch(&batch, error) || !writer.Write(&batch, error)) { + *error = "cluster: " + *error; return false; } - pending = 0; return true; }; for (uint64_t row = 0; row < assignment.root.size(); ++row) { const uint32_t root = assignment.root[static_cast(row)]; const uint64_t size = assignment.size[root]; if (size < options.min_size) continue; - arrow::Status status; - if (datasets) status &= dataset.Append(store.DatasetName(store.DatasetOf(row))); - status &= unique_id.Append(store.ids().Get(row)); - status &= cluster_id.Append(QualifiedId(store, root)); - status &= cluster_size.Append(size); - if (!status.ok()) { - *error = "cluster: building a row: " + status.message(); - return false; - } + if (datasets) + builder.AppendString(dataset, store.DatasetName(store.DatasetOf(row))); + builder.AppendString(unique_id, store.ids().Get(row)); + builder.AppendString(cluster_id, QualifiedId(store, root)); + builder.AppendUInt64(cluster_size, size); ++*written; - if (++pending >= kRowGroup && !flush()) return false; + if (builder.Rows() >= kRowGroup && !flush()) return false; } if (!flush()) return false; - const arrow::Status closed = writer->Close(); - if (!closed.ok()) { - *error = "cluster: closing '" + options.out_path + "': " + closed.message(); + if (!writer.Close(error)) { + *error = "cluster: " + *error; return false; } return true; diff --git a/src/cpplink/cluster.hpp b/src/cpplink/cluster.hpp index e7538ff..fb0749f 100644 --- a/src/cpplink/cluster.hpp +++ b/src/cpplink/cluster.hpp @@ -9,6 +9,7 @@ #include #include +#include "cpplink/arrow_c.hpp" #include "cpplink/recall.hpp" #include "cpplink/record_store.hpp" @@ -51,11 +52,18 @@ struct ClusterAssignment { bool SameCluster(uint32_t a, uint32_t b) const { return root[a] == root[b]; } }; +struct EdgeTable; + struct ClusterOptions { // Either the directory of `shard-*.bin` a run wrote, or the single `.csv` or // `.parquet` file it was merged into. A merged file names records by // `unique_id` rather than by row, so clustering one costs an id index. std::string edge_path; + // Or the run's predictions still in memory, naming rows; no index needed. + const EdgeTable* edges = nullptr; + // Or a prediction table arriving as a C Data stream -- a data frame -- which + // names records by `unique_id` and is read as a merged file is. Released. + ArrowArrayStream* stream = nullptr; // Empty writes no file; the extension picks csv or parquet, as `predict // --out` does. std::string out_path; diff --git a/src/cpplink/init.cpp b/src/cpplink/init.cpp index 4370347..b89781d 100644 --- a/src/cpplink/init.cpp +++ b/src/cpplink/init.cpp @@ -484,23 +484,39 @@ bool DraftSchema(const std::vector& paths, const DraftOptions& opti *error = "no parquet file to draft a schema from"; return false; } - const std::string& path = paths.front(); - std::vector file_columns; - if (!ReadFileColumns(path, &file_columns, error)) return false; + std::vector inputs; + for (const std::string& path : paths) { + DescribedInput input; + input.name = path; + if (!ReadFileColumns(path, &input.columns, error)) return false; + inputs.push_back(std::move(input)); + } + return DraftSchemaFrom(inputs, options, report, error); +} + +bool DraftSchemaFrom(const std::vector& inputs, + const DraftOptions& options, DraftReport* report, + std::string* error) { + if (inputs.empty()) { + *error = "no input to draft a schema from"; + return false; + } + const std::string& path = inputs.front().name; + const std::vector& file_columns = inputs.front().columns; if (file_columns.empty()) { *error = path + " has no columns"; return false; } - // The draft is from the first file; the others have to be able to run it. A + // The draft is from the first input; the others have to be able to run it. A // column the first holds must be in each of them at a type that reads the same - // way, or the schema would fail at load with a message about the second file - // that init could have given now. Columns only a later file holds are left + // way, or the schema would fail at load with a message about the second input + // that init could have given now. Columns only a later input holds are left // out, and said so. std::vector later_only; - for (size_t i = 1; i < paths.size(); ++i) { - std::vector other; - if (!ReadFileColumns(paths[i], &other, error)) return false; + for (size_t i = 1; i < inputs.size(); ++i) { + const std::vector& other = inputs[i].columns; + const std::string& other_name = inputs[i].name; for (const FileColumn& column : file_columns) { const FileColumn* match = nullptr; for (const FileColumn& candidate : other) { @@ -508,15 +524,15 @@ bool DraftSchema(const std::vector& paths, const DraftOptions& opti } if (match == nullptr) { *error = "column \"" + column.name + "\" is in " + path + " but not in " + - paths[i] + + other_name + "; a link needs the same columns " "in every input"; return false; } if (column.readable && (!match->readable || match->type != column.type)) { *error = "column \"" + column.name + "\" is " + column.arrow_type + - " in " + path + " and " + match->arrow_type + " in " + paths[i] + - "; cast one so both read as the same type"; + " in " + path + " and " + match->arrow_type + " in " + + other_name + "; cast one so both read as the same type"; return false; } } @@ -525,7 +541,7 @@ bool DraftSchema(const std::vector& paths, const DraftOptions& opti for (const FileColumn& column : file_columns) { in_first = in_first || column.name == candidate.name; } - if (!in_first) later_only.push_back(candidate.name + " (" + paths[i] + ")"); + if (!in_first) later_only.push_back(candidate.name + " (" + other_name + ")"); } } diff --git a/src/cpplink/init.hpp b/src/cpplink/init.hpp index 9e4278a..fcca1c9 100644 --- a/src/cpplink/init.hpp +++ b/src/cpplink/init.hpp @@ -96,6 +96,18 @@ bool DraftSchema(const std::string& path, const DraftOptions& options, bool DraftSchema(const std::vector& paths, const DraftOptions& options, DraftReport* report, std::string* error); +// One input as its columns describe it, from a footer or from a frame's schema. +struct DescribedInput { + std::string name; // the path, or the dataset name a frame was given + std::vector columns; +}; + +// The draft over inputs already described, which is what a frame's schema +// gives with no file to open. The two file forms read the footers and come here. +bool DraftSchemaFrom(const std::vector& inputs, + const DraftOptions& options, DraftReport* report, + std::string* error); + // Every draft is re-parsed before it is returned, so what comes out is a file // the other commands accept, not a fragment. std::string SchemaToJson(const Schema& schema); diff --git a/src/cpplink/merge_edges.cpp b/src/cpplink/merge_edges.cpp index 3395af1..f8de3e0 100644 --- a/src/cpplink/merge_edges.cpp +++ b/src/cpplink/merge_edges.cpp @@ -17,11 +17,10 @@ #include #include -#include -#include -#include - +#include "cpplink/arrow_export.hpp" +#include "cpplink/batch_loader.hpp" #include "cpplink/format.hpp" +#include "cpplink/parquet_io.hpp" #include "cpplink/predict.hpp" #include "cpplink/score.hpp" @@ -110,72 +109,47 @@ class CsvSink : public EdgeSink { std::string buffer_; }; -// The csv columns, typed, in the same order and under the same names. -std::shared_ptr MergedArrowSchema(bool datasets) { - std::vector> fields; - if (datasets) fields.push_back(arrow::field("dataset_a", arrow::utf8())); - fields.push_back(arrow::field("id_a", arrow::utf8())); - if (datasets) fields.push_back(arrow::field("dataset_b", arrow::utf8())); - fields.push_back(arrow::field("id_b", arrow::utf8())); - fields.push_back(arrow::field("gamma", arrow::uint32())); - fields.push_back(arrow::field("match_weight", arrow::float64())); - fields.push_back(arrow::field("match_probability", arrow::float64())); - return arrow::schema(fields); -} - +// The csv columns, typed, in the same order and under the same names. The rows +// are built as C Data batches and handed to the parquet writer, which is the +// same builder a data frame is handed in Python. class ParquetSink : public EdgeSink { public: ParquetSink(size_t batch_rows, bool datasets) - : batch_rows_(batch_rows == 0 ? 1 : batch_rows), - datasets_(datasets), - schema_(MergedArrowSchema(datasets)) {} + : batch_rows_(batch_rows == 0 ? 1 : batch_rows), datasets_(datasets) { + if (datasets) dataset_a_ = builder_.AddColumn("dataset_a", ExportType::kString); + id_a_ = builder_.AddColumn("id_a", ExportType::kString); + if (datasets) dataset_b_ = builder_.AddColumn("dataset_b", ExportType::kString); + id_b_ = builder_.AddColumn("id_b", ExportType::kString); + gamma_ = builder_.AddColumn("gamma", ExportType::kUInt32); + weight_ = builder_.AddColumn("match_weight", ExportType::kDouble); + probability_ = builder_.AddColumn("match_probability", ExportType::kDouble); + } bool Open(const std::string& path, std::string* error) { - path_ = path; - auto sink = arrow::io::FileOutputStream::Open(path); - if (!sink.ok()) { - *error = "merging predictions: cannot create " + path + ": " + - sink.status().message(); - return false; - } - auto props = parquet::WriterProperties::Builder() - .compression(parquet::Compression::SNAPPY) - ->build(); - auto writer = parquet::arrow::FileWriter::Open( - *schema_, arrow::default_memory_pool(), *sink, props); - if (!writer.ok()) { - *error = "merging predictions: cannot open parquet writer for " + path + - ": " + writer.status().message(); - return false; - } - writer_ = std::move(*writer); - return true; + ArrowSchema schema; + builder_.ExportSchema(&schema); + const bool ok = writer_.Open(path, schema, error); + schema.release(&schema); + if (!ok) *error = "merging predictions: " + *error; + return ok; } bool Write(const EdgeRow& edge, std::string* error) override { - arrow::Status status = id_a_.Append(edge.id_a); - status &= id_b_.Append(edge.id_b); - if (datasets_) { - status &= dataset_a_.Append(edge.dataset_a); - status &= dataset_b_.Append(edge.dataset_b); - } - status &= gamma_.Append(edge.gamma); - status &= weight_.Append(edge.weight); - status &= probability_.Append(ProbabilityForWeight(edge.weight)); - if (!status.ok()) { - *error = "merging predictions: building a row: " + status.message(); - return false; - } - ++pending_; - if (pending_ >= batch_rows_) return Flush(error); + if (datasets_) builder_.AppendString(dataset_a_, edge.dataset_a); + builder_.AppendString(id_a_, edge.id_a); + if (datasets_) builder_.AppendString(dataset_b_, edge.dataset_b); + builder_.AppendString(id_b_, edge.id_b); + builder_.AppendUInt32(gamma_, edge.gamma); + builder_.AppendDouble(weight_, edge.weight); + builder_.AppendDouble(probability_, ProbabilityForWeight(edge.weight)); + if (builder_.Rows() >= static_cast(batch_rows_)) return Flush(error); return true; } bool Close(std::string* error) override { if (!Flush(error)) return false; - const arrow::Status closed = writer_->Close(); - if (!closed.ok()) { - *error = "merging predictions: closing " + path_ + ": " + closed.message(); + if (!writer_.Close(error)) { + *error = "merging predictions: " + *error; return false; } return true; @@ -183,47 +157,21 @@ class ParquetSink : public EdgeSink { private: bool Flush(std::string* error) { - if (pending_ == 0) return true; - std::vector> arrays; - arrays.reserve(7); - arrow::Status status; - auto take = [&](arrow::ArrayBuilder* builder) { - std::shared_ptr array; - status &= builder->Finish(&array); - arrays.push_back(std::move(array)); - }; - if (datasets_) take(&dataset_a_); - take(&id_a_); - if (datasets_) take(&dataset_b_); - take(&id_b_); - take(&gamma_); - take(&weight_); - take(&probability_); - if (!status.ok()) { - *error = "merging predictions: finishing a batch: " + status.message(); + if (builder_.Rows() == 0) return true; + ArrowArray batch; + if (!builder_.ExportBatch(&batch, error) || !writer_.Write(&batch, error)) { + *error = "merging predictions: " + *error; return false; } - const auto table = - arrow::Table::Make(schema_, arrays, static_cast(pending_)); - status = writer_->WriteTable(*table, static_cast(pending_)); - if (!status.ok()) { - *error = "merging predictions: writing a row group to " + path_ + ": " + - status.message(); - return false; - } - pending_ = 0; return true; } size_t batch_rows_; bool datasets_; - std::shared_ptr schema_; - arrow::StringBuilder dataset_a_, dataset_b_, id_a_, id_b_; - arrow::UInt32Builder gamma_; - arrow::DoubleBuilder weight_, probability_; - size_t pending_ = 0; - std::string path_; - std::unique_ptr writer_; + BatchBuilder builder_; + int dataset_a_ = -1, id_a_ = -1, dataset_b_ = -1, id_b_ = -1; + int gamma_ = -1, weight_ = -1, probability_ = -1; + ParquetWriter writer_; }; bool CollectShards(const std::string& dir, MergeSource wanted, @@ -527,6 +475,142 @@ bool MergeStagedShards(const RecordStore& store, const std::string& staging, return true; } +bool ReadPredictionStream(ArrowArrayStream* stream, const std::string& what, + const std::string& source, + const std::vector& required, EdgeColumns* columns, + const EdgeVisitor& visit, std::string* error) { + auto release = [&]() { + if (stream->release != nullptr) stream->release(stream); + }; + ArrowSchema schema; + schema.release = nullptr; + if (const int code = stream->get_schema(stream, &schema); code != 0) { + const char* detail = + stream->get_last_error == nullptr ? nullptr : stream->get_last_error(stream); + *error = what + ": cannot read the schema of '" + source + + "': " + (detail == nullptr ? "error " + std::to_string(code) : detail); + release(); + return false; + } + auto fail = [&](const std::string& message) { + *error = what + ": " + message; + schema.release(&schema); + release(); + return false; + }; + + const int id_a = FieldIndex(schema, "id_a"); + const int id_b = FieldIndex(schema, "id_b"); + if (id_a < 0 || id_b < 0) { + return fail("'" + source + "' has no column \"" + (id_a < 0 ? "id_a" : "id_b") + + "\", so it is not a cpplink prediction table"); + } + const int dataset_a = FieldIndex(schema, "dataset_a"); + const int dataset_b = FieldIndex(schema, "dataset_b"); + if ((dataset_a < 0) != (dataset_b < 0)) { + return fail("'" + source + "' has " + + (dataset_a < 0 ? "dataset_b" : "dataset_a") + " but no " + + (dataset_a < 0 ? "dataset_a" : "dataset_b")); + } + const int weight = FieldIndex(schema, "match_weight"); + const int gamma = FieldIndex(schema, "gamma"); + for (const std::string& name : required) { + if (FieldIndex(schema, name) < 0) { + return fail("'" + source + "' has no column \"" + name + + "\", so it is not a cpplink prediction table"); + } + } + columns->datasets = dataset_a >= 0; + columns->weight = weight >= 0; + columns->gamma = gamma >= 0; + + while (true) { + ArrowArray batch; + batch.release = nullptr; + if (const int code = stream->get_next(stream, &batch); code != 0) { + const char* detail = stream->get_last_error == nullptr + ? nullptr + : stream->get_last_error(stream); + return fail("cannot read the next batch of '" + source + "': " + + (detail == nullptr ? "error " + std::to_string(code) : detail)); + } + if (batch.release == nullptr) break; + TextReader a, b, set_a, set_b; + NumberReader weights, gammas; + std::string bind_error; + const bool bound = + a.Bind(schema, batch, id_a, &bind_error) && + b.Bind(schema, batch, id_b, &bind_error) && + (!columns->datasets || (set_a.Bind(schema, batch, dataset_a, &bind_error) && + set_b.Bind(schema, batch, dataset_b, &bind_error))) && + (!columns->weight || weights.Bind(schema, batch, weight, &bind_error)) && + (!columns->gamma || gammas.Bind(schema, batch, gamma, &bind_error)); + if (!bound) { + batch.release(&batch); + return fail("'" + source + + "' holds a column in a type a cpplink prediction " + + "table does not use: " + bind_error); + } + char scratch_a[24], scratch_b[24], scratch_c[24], scratch_d[24]; + for (int64_t row = 0; row < batch.length; ++row) { + EdgeRow edge; + if (!a.At(row, &edge.id_a, &scratch_a)) edge.id_a = std::string_view(); + if (!b.At(row, &edge.id_b, &scratch_b)) edge.id_b = std::string_view(); + if (columns->datasets) { + if (!set_a.At(row, &edge.dataset_a, &scratch_c)) { + edge.dataset_a = std::string_view(); + } + if (!set_b.At(row, &edge.dataset_b, &scratch_d)) { + edge.dataset_b = std::string_view(); + } + } + double value = 0.0; + if (columns->weight && weights.At(row, &value)) edge.weight = value; + if (columns->gamma && gammas.At(row, &value)) { + edge.gamma = static_cast(value); + } + if (!visit(edge, error)) { + batch.release(&batch); + schema.release(&schema); + release(); + return false; + } + } + batch.release(&batch); + } + schema.release(&schema); + release(); + return true; +} + +bool ReadPredictionFile(const std::string& path, const std::string& what, + const std::vector& required, EdgeColumns* columns, + const EdgeVisitor& visit, std::string* error) { + // Only the columns the shape has are read: the pattern and the posterior + // stay in the file for whoever else opens it. + ArrowSchema schema; + schema.release = nullptr; + if (!ReadParquetSchema(path, &schema, error)) { + *error = what + ": " + *error; + return false; + } + std::vector names; + for (const char* name : {"id_a", "id_b", "dataset_a", "dataset_b"}) { + if (FieldIndex(schema, name) >= 0) names.push_back(name); + } + for (const std::string& name : required) { + if (FieldIndex(schema, name) >= 0) names.push_back(name); + } + schema.release(&schema); + ArrowArrayStream stream; + stream.release = nullptr; + if (!OpenParquetStream(path, names, &stream, error)) { + *error = what + ": " + *error; + return false; + } + return ReadPredictionStream(&stream, what, path, required, columns, visit, error); +} + bool MergedFormatOf(const std::string& path, MergeFormat* format) { const std::string suffix = std::filesystem::path(path).extension().string(); if (suffix == ".parquet" || suffix == ".pq") { diff --git a/src/cpplink/merge_edges.hpp b/src/cpplink/merge_edges.hpp index 27d9f71..bebf296 100644 --- a/src/cpplink/merge_edges.hpp +++ b/src/cpplink/merge_edges.hpp @@ -4,12 +4,14 @@ #pragma once #include +#include #include #include #include #include #include +#include "cpplink/arrow_c.hpp" #include "cpplink/record_store.hpp" namespace cpplink { @@ -76,6 +78,35 @@ struct EdgeRow { double weight = 0.0; }; +// Which columns a prediction file's rows were read with. The ids are always +// there; the datasets come together or not at all; the weight and the pattern +// are each read where the file holds them, and a caller that needs one asks +// for it by name and is refused before any row where it is missing. +struct EdgeColumns { + bool datasets = false; + bool weight = false; + bool gamma = false; +}; + +using EdgeVisitor = std::function; + +// Every row of a prediction table read through the C Data Interface, whether it +// came from a merged parquet file or from a data frame: the same reader either +// way, so `cluster --predictions x.parquet` and `cluster(predictions_df)` cannot +// read a row differently. `required` names the columns beyond the ids the caller +// cannot do without, of `match_weight` and `gamma`; `what` names the caller in +// the message. A row with a null id or dataset is passed with that view empty. +// The stream is released on return. +bool ReadPredictionStream(ArrowArrayStream* stream, const std::string& what, + const std::string& source, + const std::vector& required, EdgeColumns* columns, + const EdgeVisitor& visit, std::string* error); + +// The same over a merged parquet file, reading only the columns the shape has. +bool ReadPredictionFile(const std::string& path, const std::string& what, + const std::vector& required, EdgeColumns* columns, + const EdgeVisitor& visit, std::string* error); + // Which of the two csv shapes a file is in, read off its header. struct EdgeCsvLayout { bool datasets = false; // `dataset_a,id_a,dataset_b,id_b,...` diff --git a/src/cpplink/parquet_io.hpp b/src/cpplink/parquet_io.hpp new file mode 100644 index 0000000..f0895ec --- /dev/null +++ b/src/cpplink/parquet_io.hpp @@ -0,0 +1,52 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +#pragma once + +#include +#include +#include + +#include "cpplink/arrow_c.hpp" + +namespace cpplink { + +// The whole of what cpplink asks of Arrow C++: a parquet file read as a stream +// of C Data batches, and a stream of C Data batches written as a parquet file. +// Everything above this line reads and writes the C structs, so the parquet +// library is a leaf that `CPPLINK_WITH_ARROW` swaps for a stub, and a wheel that +// takes its input from a data frame links no Arrow at all. The stub fails every +// call with the reason. + +// Whether this build can read and write parquet. +bool ParquetSupported(); + +// The file's schema, as the C struct. The caller releases it. +bool ReadParquetSchema(const std::string& path, ArrowSchema* out, std::string* error); + +// The file as a stream of batches, one row group at a time, holding only the +// named columns (all of them where `columns` is empty). A name the file lacks +// is an error here rather than a missing child later. The caller releases the +// stream, which closes the file. +bool OpenParquetStream(const std::string& path, const std::vector& columns, + ArrowArrayStream* out, std::string* error); + +// A parquet file written batch by batch; each batch is one row group. +class ParquetWriter { + public: + ParquetWriter(); + ~ParquetWriter(); + + // Creates the file for batches of this schema. The schema is read, not + // taken: the caller still releases it. + bool Open(const std::string& path, const ArrowSchema& schema, std::string* error); + // Writes and releases the batch, whichever way it goes. + bool Write(ArrowArray* batch, std::string* error); + bool Close(std::string* error); + + private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace cpplink diff --git a/src/cpplink/parquet_io_arrow.cpp b/src/cpplink/parquet_io_arrow.cpp new file mode 100644 index 0000000..5d15fd7 --- /dev/null +++ b/src/cpplink/parquet_io_arrow.cpp @@ -0,0 +1,214 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +// The Arrow C++ side of `parquet_io.hpp`, and the only file in the library that +// includes an Arrow header. Compiled under CPPLINK_WITH_ARROW=ON; the stub takes +// its place otherwise. + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "cpplink/parquet_io.hpp" + +namespace cpplink { +namespace { + +bool OpenReader(const std::string& path, + std::unique_ptr* reader, + std::shared_ptr* schema, std::string* error) { + auto input = arrow::io::ReadableFile::Open(path); + if (!input.ok()) { + *error = "cannot open " + path + ": " + input.status().message(); + return false; + } + parquet::arrow::FileReaderBuilder builder; + arrow::Status status = builder.Open(*input); + if (!status.ok()) { + *error = "cannot read parquet metadata: " + status.message(); + return false; + } + auto built = builder.Build(); + if (!built.ok()) { + *error = "cannot open parquet reader: " + built.status().message(); + return false; + } + *reader = std::move(*built); + status = (*reader)->GetSchema(schema); + if (!status.ok()) { + *error = "cannot read parquet schema: " + status.message(); + return false; + } + return true; +} + +} // namespace + +bool ParquetSupported() { return true; } + +bool ReadParquetSchema(const std::string& path, ArrowSchema* out, std::string* error) { + std::unique_ptr reader; + std::shared_ptr schema; + if (!OpenReader(path, &reader, &schema, error)) return false; + const arrow::Status status = arrow::ExportSchema(*schema, out); + if (!status.ok()) { + *error = "cannot export parquet schema: " + status.message(); + return false; + } + return true; +} + +bool OpenParquetStream(const std::string& path, const std::vector& columns, + ArrowArrayStream* out, std::string* error) { + std::unique_ptr reader; + std::shared_ptr schema; + if (!OpenReader(path, &reader, &schema, error)) return false; + std::vector indices; + for (const std::string& name : columns) { + const int index = schema->GetFieldIndex(name); + if (index < 0) { + *error = "column \"" + name + "\" is not in " + path; + return false; + } + indices.push_back(index); + } + if (columns.empty()) { + for (int i = 0; i < schema->num_fields(); ++i) indices.push_back(i); + } + // One row group at a time, read whole and handed out as its batches, so a + // group's buffers are released before the next is read and the two + // representations are never both resident for the whole file. The reader + // owns the file, so releasing the stream closes it. + class RowGroups : public arrow::RecordBatchReader { + public: + RowGroups(std::unique_ptr file, + std::shared_ptr schema, std::vector indices) + : file_(std::move(file)), + schema_(std::move(schema)), + indices_(std::move(indices)) {} + std::shared_ptr schema() const override { return schema_; } + arrow::Status ReadNext(std::shared_ptr* batch) override { + while (true) { + if (batches_) { + ARROW_RETURN_NOT_OK(batches_->ReadNext(batch)); + if (*batch != nullptr) return arrow::Status::OK(); + batches_.reset(); + table_.reset(); + } + if (group_ >= file_->num_row_groups()) { + batch->reset(); + return arrow::Status::OK(); + } + ARROW_ASSIGN_OR_RAISE(table_, file_->ReadRowGroup(group_++, indices_)); + batches_ = std::make_unique(*table_); + } + } + + private: + std::unique_ptr file_; + std::shared_ptr schema_; + std::vector indices_; + int group_ = 0; + std::shared_ptr table_; + std::unique_ptr batches_; + }; + // The schema of the columns read, in the order they were asked for. + std::vector> fields; + for (const int index : indices) fields.push_back(schema->field(index)); + auto rows = std::make_shared(std::move(reader), arrow::schema(fields), + std::move(indices)); + const arrow::Status status = arrow::ExportRecordBatchReader(std::move(rows), out); + if (!status.ok()) { + *error = "cannot export " + path + ": " + status.message(); + return false; + } + return true; +} + +struct ParquetWriter::Impl { + std::string path; + std::shared_ptr schema; + std::unique_ptr writer; +}; + +ParquetWriter::ParquetWriter() : impl_(std::make_unique()) {} +ParquetWriter::~ParquetWriter() = default; + +bool ParquetWriter::Open(const std::string& path, const ArrowSchema& schema, + std::string* error) { + impl_->path = path; + // Importing consumes the struct, so a copy is imported and the caller's is + // left for the caller to release. + ArrowSchema copy = schema; + copy.release = [](ArrowSchema* s) { s->release = nullptr; }; + auto imported = arrow::ImportSchema(©); + if (!imported.ok()) { + *error = + "cannot import the schema for " + path + ": " + imported.status().message(); + return false; + } + impl_->schema = std::move(*imported); + auto sink = arrow::io::FileOutputStream::Open(path); + if (!sink.ok()) { + *error = "cannot create " + path + ": " + sink.status().message(); + return false; + } + auto props = parquet::WriterProperties::Builder() + .compression(parquet::Compression::SNAPPY) + ->build(); + auto writer = parquet::arrow::FileWriter::Open( + *impl_->schema, arrow::default_memory_pool(), *sink, props); + if (!writer.ok()) { + *error = + "cannot open parquet writer for " + path + ": " + writer.status().message(); + return false; + } + impl_->writer = std::move(*writer); + return true; +} + +bool ParquetWriter::Write(ArrowArray* batch, std::string* error) { + auto imported = arrow::ImportRecordBatch(batch, impl_->schema); + if (!imported.ok()) { + if (batch->release != nullptr) batch->release(batch); + *error = "cannot import a batch for " + impl_->path + ": " + + imported.status().message(); + return false; + } + const auto& rows = *imported; + if (rows->num_rows() == 0) return true; + // Each batch is one row group: `WriteTable` closes the group where + // `WriteRecordBatch` would buffer batches into one until a million rows. + auto table = arrow::Table::FromRecordBatches(impl_->schema, {rows}); + if (!table.ok()) { + *error = + "cannot table a batch for " + impl_->path + ": " + table.status().message(); + return false; + } + const arrow::Status status = impl_->writer->WriteTable(**table, rows->num_rows()); + if (!status.ok()) { + *error = "writing a row group to " + impl_->path + ": " + status.message(); + return false; + } + return true; +} + +bool ParquetWriter::Close(std::string* error) { + if (!impl_->writer) return true; + const arrow::Status closed = impl_->writer->Close(); + impl_->writer.reset(); + if (!closed.ok()) { + *error = "closing " + impl_->path + ": " + closed.message(); + return false; + } + return true; +} + +} // namespace cpplink diff --git a/src/cpplink/parquet_io_stub.cpp b/src/cpplink/parquet_io_stub.cpp new file mode 100644 index 0000000..4a8d100 --- /dev/null +++ b/src/cpplink/parquet_io_stub.cpp @@ -0,0 +1,55 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +// `parquet_io.hpp` without Arrow: every call fails and says why. Compiled under +// CPPLINK_WITH_ARROW=OFF, which is how the Python wheel is built -- its input is +// a data frame and its output a data frame, and neither needs a file. + +#include +#include + +#include "cpplink/parquet_io.hpp" + +namespace cpplink { +namespace { + +const char* kUnsupported = + "this build of cpplink cannot read or write parquet: it was built without Arrow"; + +} // namespace + +bool ParquetSupported() { return false; } + +bool ReadParquetSchema(const std::string&, ArrowSchema*, std::string* error) { + *error = kUnsupported; + return false; +} + +bool OpenParquetStream(const std::string&, const std::vector&, + ArrowArrayStream*, std::string* error) { + *error = kUnsupported; + return false; +} + +struct ParquetWriter::Impl {}; + +ParquetWriter::ParquetWriter() = default; +ParquetWriter::~ParquetWriter() = default; + +bool ParquetWriter::Open(const std::string&, const ArrowSchema&, std::string* error) { + *error = kUnsupported; + return false; +} + +bool ParquetWriter::Write(ArrowArray* batch, std::string* error) { + if (batch->release != nullptr) batch->release(batch); + *error = kUnsupported; + return false; +} + +bool ParquetWriter::Close(std::string* error) { + *error = kUnsupported; + return false; +} + +} // namespace cpplink diff --git a/src/cpplink/parquet_loader.cpp b/src/cpplink/parquet_loader.cpp index 28dd103..45188cb 100644 --- a/src/cpplink/parquet_loader.cpp +++ b/src/cpplink/parquet_loader.cpp @@ -4,476 +4,51 @@ #include "cpplink/parquet_loader.hpp" #include -#include -#include #include -#include -#include #include -#include #include #include -#include -#include -#include +#include "cpplink/batch_loader.hpp" +#include "cpplink/parquet_io.hpp" namespace cpplink { namespace { -// Calls fn on the array as its concrete integer type, or returns false when the -// column is not an integer of any width or sign. -template -bool WithIntegerArray(const arrow::Array& array, Fn&& fn) { - switch (array.type_id()) { - case arrow::Type::INT8: - fn(static_cast(array)); - return true; - case arrow::Type::INT16: - fn(static_cast(array)); - return true; - case arrow::Type::INT32: - fn(static_cast(array)); - return true; - case arrow::Type::INT64: - fn(static_cast(array)); - return true; - case arrow::Type::UINT8: - fn(static_cast(array)); - return true; - case arrow::Type::UINT16: - fn(static_cast(array)); - return true; - case arrow::Type::UINT32: - fn(static_cast(array)); - return true; - case arrow::Type::UINT64: - fn(static_cast(array)); - return true; - default: - return false; - } -} - -// Visits every value of a text-like column in order as `visit(value, is_null)`. -// A file written from Python rarely holds the types the schema names: pandas and -// polars write `large_string` for text, and a postcode, a year or an id that was -// numeric in the frame arrives as an integer. Both are the string the user meant, -// so an integer is visited as its decimal digits and interned like any other -// value, which is what makes "2000" in one file and 2000 in another the same id. -template -void VisitStrings(const arrow::Array& array, Visit& visit) { - const auto& typed = static_cast(array); - for (int64_t i = 0; i < typed.length(); ++i) { - if (typed.IsNull(i)) { - visit(std::string_view(), true); - } else { - visit(std::string_view(typed.GetView(i)), false); - } - } -} - -template -bool ForEachText(const arrow::Array& array, Visit&& visit, std::string* error) { - if (array.type_id() == arrow::Type::STRING) { - VisitStrings(array, visit); - return true; +// The exported struct releases itself through the callback the producer put in +// it; this guard makes sure it does whichever way a function leaves. +struct SchemaGuard { + ArrowSchema schema; + SchemaGuard() { schema.release = nullptr; } + ~SchemaGuard() { + if (schema.release != nullptr) schema.release(&schema); } - if (array.type_id() == arrow::Type::LARGE_STRING) { - VisitStrings(array, visit); - return true; - } - const bool integer = WithIntegerArray(array, [&](const auto& typed) { - char digits[24]; // -9223372036854775808 is 20 characters - for (int64_t i = 0; i < typed.length(); ++i) { - if (typed.IsNull(i)) { - visit(std::string_view(), true); - continue; - } - const auto result = - std::to_chars(digits, digits + sizeof(digits), typed.Value(i)); - visit(std::string_view(digits, result.ptr - digits), false); - } - }); - if (integer) return true; - *error = "expected a string or integer column, found " + array.type()->ToString(); - return false; -} - -// Days since the Unix epoch, whatever the column's temporal unit. -bool AppendDates(const arrow::Array& array, DateColumn* column, std::string* error) { - switch (array.type_id()) { - case arrow::Type::DATE32: { - const auto& typed = static_cast(array); - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back(typed.IsNull(i) ? kNullDate : typed.Value(i)); - } - return true; - } - case arrow::Type::DATE64: { - const auto& typed = static_cast(array); - constexpr int64_t kMillisPerDay = 86400000; - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back( - typed.IsNull(i) - ? kNullDate - : static_cast(typed.Value(i) / kMillisPerDay)); - } - return true; - } - case arrow::Type::TIMESTAMP: { - const auto& typed = static_cast(array); - const auto& type = static_cast(*typed.type()); - int64_t per_day = 86400; - switch (type.unit()) { - case arrow::TimeUnit::SECOND: - per_day = 86400LL; - break; - case arrow::TimeUnit::MILLI: - per_day = 86400000LL; - break; - case arrow::TimeUnit::MICRO: - per_day = 86400000000LL; - break; - case arrow::TimeUnit::NANO: - per_day = 86400000000000LL; - break; - } - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back( - typed.IsNull(i) ? kNullDate - : static_cast(typed.Value(i) / per_day)); - } - return true; - } - default: - *error = - "expected a date or timestamp column, found " + array.type()->ToString(); - return false; - } -} - -bool AppendDoubles(const arrow::Array& array, DoubleColumn* column, std::string* error) { - const double kMissing = std::numeric_limits::quiet_NaN(); - if (array.type_id() == arrow::Type::DOUBLE) { - const auto& typed = static_cast(array); - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back(typed.IsNull(i) ? kMissing : typed.Value(i)); - } - return true; - } - if (array.type_id() == arrow::Type::FLOAT) { - const auto& typed = static_cast(array); - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back(typed.IsNull(i) ? kMissing : typed.Value(i)); - } - return true; - } - *error = "expected a floating point column, found " + array.type()->ToString(); - return false; -} +}; -bool AppendBooleans(const arrow::Array& array, BooleanColumn* column, - std::string* error) { - if (array.type_id() == arrow::Type::BOOL) { - const auto& typed = static_cast(array); - for (int64_t i = 0; i < typed.length(); ++i) { - column->values.push_back(typed.IsNull(i) ? kNullBoolean - : typed.Value(i) ? int8_t{1} - : int8_t{0}); - } - return true; - } - *error = "expected a boolean column, found " + array.type()->ToString(); - return false; -} - -bool AppendStrings(const arrow::Array& array, StringColumn* column, std::string* error) { - return ForEachText( - array, - [&](std::string_view value, bool is_null) { - column->ids.push_back(is_null ? kNullId : column->dict.Intern(value)); - }, - error); -} - -// Slices the interned elements of a list chunk into one sorted, deduplicated cell -// per row, whether the offsets are 32-bit (`list`) or 64-bit (`large_list`). -template -void AppendListRows(const ListArray& lists, const std::vector& elements, - StringListColumn* column) { - std::vector row; - for (int64_t i = 0; i < lists.length(); ++i) { - row.clear(); - if (!lists.IsNull(i)) { - const int64_t begin = lists.value_offset(i); - const int64_t end = lists.value_offset(i + 1); - for (int64_t j = begin; j < end; ++j) { - if (elements[j] != kNullId) row.push_back(elements[j]); - } - } - // Sorted and deduplicated here so overlap is a linear merge later. - std::sort(row.begin(), row.end()); - row.erase(std::unique(row.begin(), row.end()), row.end()); - column->ids.insert(column->ids.end(), row.begin(), row.end()); - column->offsets.push_back(column->ids.size()); - } -} - -bool AppendStringLists(const arrow::Array& array, StringListColumn* column, - std::string* error) { - const bool large = array.type_id() == arrow::Type::LARGE_LIST; - if (array.type_id() != arrow::Type::LIST && !large) { - *error = "expected a list column, found " + array.type()->ToString(); - return false; - } - // The elements are interned first, as one flat run over the child array, so a - // list of integers goes through the same visitor a scalar column does; the rows - // are then cut out of that run by the list's own offsets. - const std::shared_ptr values = - large ? static_cast(array).values() - : static_cast(array).values(); - std::vector elements; - elements.reserve(static_cast(values->length())); - std::string element_error; - const bool ok = ForEachText( - *values, - [&](std::string_view value, bool is_null) { - elements.push_back(is_null ? kNullId : column->dict.Intern(value)); - }, - &element_error); - if (!ok) { - *error = "expected a list of strings or integers, found a list of " + - values->type()->ToString(); - return false; - } - if (column->offsets.empty()) column->offsets.push_back(0); - if (large) { - AppendListRows(static_cast(array), elements, - column); - } else { - AppendListRows(static_cast(array), elements, column); - } - return true; -} - -bool AppendIds(const arrow::Array& array, IdColumn* column, std::string* error) { - std::string text_error; - const bool ok = ForEachText( - array, [&](std::string_view value, bool) { column->Append(value); }, &text_error); - if (ok) return true; - *error = "expected a string or integer unique_id column, found " + - array.type()->ToString(); - return false; -} - -bool AppendColumn(const arrow::ChunkedArray& chunked, ColumnType type, Column* column, - std::string* error) { - for (const auto& chunk : chunked.chunks()) { - bool ok = false; - switch (type) { - case ColumnType::kString: - ok = AppendStrings(*chunk, &std::get(*column), error); - break; - case ColumnType::kStringList: - ok = AppendStringLists(*chunk, &std::get(*column), - error); - break; - case ColumnType::kDate: - ok = AppendDates(*chunk, &std::get(*column), error); - break; - case ColumnType::kBoolean: - ok = AppendBooleans(*chunk, &std::get(*column), error); - break; - case ColumnType::kDouble: - ok = AppendDoubles(*chunk, &std::get(*column), error); - break; - } - if (!ok) return false; - } - return true; -} - -bool OpenSchema(const std::string& path, std::shared_ptr* file_schema, - std::string* error) { - auto input_result = arrow::io::ReadableFile::Open(path); - if (!input_result.ok()) { - *error = "cannot open " + path + ": " + input_result.status().message(); - return false; - } - parquet::arrow::FileReaderBuilder builder; - auto status = builder.Open(*input_result); - if (!status.ok()) { - *error = "cannot read parquet metadata: " + status.message(); - return false; - } - auto reader_result = builder.Build(); - if (!reader_result.ok()) { - *error = "cannot open parquet reader: " + reader_result.status().message(); - return false; - } - status = (*reader_result)->GetSchema(file_schema); - if (!status.ok()) { - *error = "cannot read parquet schema: " + status.message(); - return false; - } - return true; -} - -// The column type an Arrow type reads as, or false where none does. An integer -// is text: a postcode, a year or a phone number that was numeric in the frame is -// still the value it spells, and a writer who wanted a double wrote one. -bool ColumnTypeOf(const arrow::DataType& type, ColumnType* out) { - switch (type.id()) { - case arrow::Type::STRING: - case arrow::Type::LARGE_STRING: - case arrow::Type::INT8: - case arrow::Type::INT16: - case arrow::Type::INT32: - case arrow::Type::INT64: - case arrow::Type::UINT8: - case arrow::Type::UINT16: - case arrow::Type::UINT32: - case arrow::Type::UINT64: - *out = ColumnType::kString; - return true; - case arrow::Type::LIST: - case arrow::Type::LARGE_LIST: - *out = ColumnType::kStringList; - return true; - case arrow::Type::DATE32: - case arrow::Type::DATE64: - case arrow::Type::TIMESTAMP: - *out = ColumnType::kDate; - return true; - case arrow::Type::BOOL: - *out = ColumnType::kBoolean; - return true; - case arrow::Type::FLOAT: - case arrow::Type::DOUBLE: - *out = ColumnType::kDouble; - return true; - default: - return false; - } -} - -// Appends one file's rows to the store without finalizing it: several files may -// still be coming, and the dictionaries and term frequencies are only correct once -// every row that will ever be in the store has been read. -bool AppendOneFile(const std::string& path, const Schema& schema, RecordStore* store, - uint64_t* rows_read, int* row_groups_read, std::string* error) { - auto input_result = arrow::io::ReadableFile::Open(path); - if (!input_result.ok()) { - *error = "cannot open " + path + ": " + input_result.status().message(); - return false; - } - - parquet::arrow::FileReaderBuilder builder; - auto status = builder.Open(*input_result); - if (!status.ok()) { - *error = "cannot read parquet metadata: " + status.message(); - return false; - } - auto reader_result = builder.Build(); - if (!reader_result.ok()) { - *error = "cannot open parquet reader: " + reader_result.status().message(); - return false; - } - std::unique_ptr reader = std::move(*reader_result); - - std::shared_ptr file_schema; - status = reader->GetSchema(&file_schema); - if (!status.ok()) { - *error = "cannot read parquet schema: " + status.message(); - return false; - } - - // Resolve every column up front so a typo fails before any work is done. A - // derived column is not among them: it is computed from another column when - // the store is finalized, and the file is not expected to hold it. - constexpr size_t kIdTarget = static_cast(-1); - std::vector indices; +// The columns a schema reads from a file: the id and every column that is not +// derived, so a file is read for what the store holds and nothing else. +std::vector ColumnsToRead(const Schema& schema) { std::vector names; - std::vector targets; // the store column each name fills - int id_index = -1; - if (!schema.unique_id.empty()) { - id_index = file_schema->GetFieldIndex(schema.unique_id); - if (id_index < 0) { - *error = "unique_id column \"" + schema.unique_id + "\" is not in the file"; - return false; - } - indices.push_back(id_index); - names.push_back(schema.unique_id); - targets.push_back(kIdTarget); - } - for (size_t i = 0; i < schema.columns.size(); ++i) { - const ColumnSpec& spec = schema.columns[i]; - if (spec.IsDerived()) continue; - const int index = file_schema->GetFieldIndex(spec.name); - if (index < 0) { - *error = "column \"" + spec.name + "\" is not in the file"; - return false; - } - indices.push_back(index); - names.push_back(spec.name); - targets.push_back(i); + if (!schema.unique_id.empty()) names.push_back(schema.unique_id); + for (const ColumnSpec& spec : schema.columns) { + if (!spec.IsDerived()) names.push_back(spec.name); } - - const int row_groups = reader->num_row_groups(); - uint64_t rows = 0; - for (int group = 0; group < row_groups; ++group) { - std::shared_ptr table; - status = reader->ReadRowGroup(group, indices, &table); - if (!status.ok()) { - *error = "cannot read row group " + std::to_string(group) + ": " + - status.message(); - return false; - } - for (size_t i = 0; i < names.size(); ++i) { - const auto chunked = table->GetColumnByName(names[i]); - if (chunked == nullptr) { - *error = "column \"" + names[i] + "\" vanished from row group " + - std::to_string(group); - return false; - } - if (targets[i] == kIdTarget) { - for (const auto& chunk : chunked->chunks()) { - if (!AppendIds(*chunk, &store->mutable_ids(), error)) { - *error = schema.unique_id + ": " + *error; - return false; - } - } - continue; - } - const size_t column_index = targets[i]; - if (!AppendColumn(*chunked, schema.columns[column_index].type, - &store->mutable_column(column_index), error)) { - *error = names[i] + ": " + *error; - return false; - } - } - rows += static_cast(table->num_rows()); - // The row group's Arrow buffers are released here, before the next read. - } - - *rows_read = rows; - *row_groups_read = row_groups; - return true; + return names; } } // namespace bool ReadFileColumns(const std::string& path, std::vector* columns, std::string* error) { - std::shared_ptr file_schema; - if (!OpenSchema(path, &file_schema, error)) return false; + SchemaGuard file; + if (!ReadParquetSchema(path, &file.schema, error)) return false; columns->clear(); - for (const auto& field : file_schema->fields()) { + for (int64_t i = 0; i < file.schema.n_children; ++i) { + const ArrowSchema& field = *file.schema.children[i]; FileColumn column; - column.name = field->name(); - column.arrow_type = field->type()->ToString(); - column.readable = ColumnTypeOf(*field->type(), &column.type); + column.name = field.name == nullptr ? "" : field.name; + column.arrow_type = ArrowTypeName(field); + column.readable = ColumnTypeOf(field, &column.type); columns->push_back(std::move(column)); } return true; @@ -485,23 +60,9 @@ bool ResolveColumnTypes(const std::vector& paths, Schema* schema, *error = "no parquet file to read the column types from"; return false; } - std::shared_ptr file_schema; - if (!OpenSchema(paths.front(), &file_schema, error)) return false; - for (ColumnSpec& spec : schema->columns) { - if (spec.type_declared || spec.IsDerived()) continue; - const int index = file_schema->GetFieldIndex(spec.name); - if (index < 0) { - *error = "column \"" + spec.name + "\" is not in " + paths.front(); - return false; - } - const auto& type = *file_schema->field(index)->type(); - if (!ColumnTypeOf(type, &spec.type)) { - *error = "column \"" + spec.name + "\" is " + type.ToString() + " in " + - paths.front() + ", which cpplink cannot read; cast it to a string"; - return false; - } - } - return CheckTypes(schema, error); + SchemaGuard file; + if (!ReadParquetSchema(paths.front(), &file.schema, error)) return false; + return ResolveColumnTypesFrom(file.schema, paths.front(), schema, error); } std::vector DatasetNamesFor(const std::vector& paths) { @@ -524,53 +85,38 @@ std::vector DatasetNamesFor(const std::vector& paths) bool LoadParquetFiles(const std::vector& paths, const Schema& schema, RecordStore* store, LoadStats* stats, std::string* error) { - const auto started = std::chrono::steady_clock::now(); if (paths.empty()) { *error = "no input file was given"; return false; } - - // Boundaries are recorded as they are crossed, so a dataset is the row range - // between two of them and nothing per row has to be stored. - std::vector starts; - std::vector per_file; - starts.push_back(0); - uint64_t total = 0; - int row_groups = 0; - for (const std::string& path : paths) { - uint64_t rows = 0; - int groups = 0; - if (!AppendOneFile(path, schema, store, &rows, &groups, error)) { - // With several inputs the message has to say which one failed; a - // column missing from the second file reads the same as one missing - // from the first otherwise. - if (paths.size() > 1) *error = path + ": " + *error; + // Every file is opened before any is read, so a missing column in the last + // input fails before the first is loaded. Opening reads the footer only. + const std::vector columns = ColumnsToRead(schema); + std::vector streams(paths.size()); + std::vector inputs; + const std::vector names = DatasetNamesFor(paths); + for (size_t i = 0; i < paths.size(); ++i) { + streams[i].release = nullptr; + if (!OpenParquetStream(paths[i], columns, &streams[i], error)) { + for (size_t j = 0; j < i; ++j) streams[j].release(&streams[j]); + if (paths.size() > 1) *error = paths[i] + ": " + *error; return false; } - total += rows; - row_groups += groups; - per_file.push_back(rows); - starts.push_back(total); + inputs.push_back({names.empty() ? std::string() : names[i], &streams[i]}); } - - store->set_num_records(total); - // One dataset is the absence of a boundary, not a boundary at each end: the - // dedup path then pays no lookup at all. + if (LoadStreams(inputs, schema, store, stats, error)) return true; + // `LoadStreams` names the input by its dataset name; the file is what the + // user typed, so say that where there is more than one. if (paths.size() > 1) { - store->set_datasets(std::move(starts)); - store->set_dataset_names(DatasetNamesFor(paths)); - } - store->Finalize(); - - if (stats != nullptr) { - stats->rows = total; - stats->row_groups = row_groups; - stats->dataset_rows = std::move(per_file); - stats->seconds = - std::chrono::duration(std::chrono::steady_clock::now() - started) - .count(); + for (size_t i = 0; i < paths.size(); ++i) { + const std::string prefix = names[i] + ": "; + if (error->compare(0, prefix.size(), prefix) == 0) { + *error = paths[i] + ": " + error->substr(prefix.size()); + break; + } + } } - return true; + return false; } bool LoadParquet(const std::string& path, const Schema& schema, RecordStore* store, diff --git a/src/cpplink/parquet_loader.hpp b/src/cpplink/parquet_loader.hpp index 90e8692..c5c2b5d 100644 --- a/src/cpplink/parquet_loader.hpp +++ b/src/cpplink/parquet_loader.hpp @@ -6,18 +6,12 @@ #include #include +#include "cpplink/batch_loader.hpp" #include "cpplink/record_store.hpp" #include "cpplink/schema.hpp" namespace cpplink { -struct LoadStats { - uint64_t rows = 0; - int row_groups = 0; - double seconds = 0.0; - std::vector dataset_rows; // rows contributed by each input, in order -}; - // One column of a parquet file as its footer describes it: the name, the Arrow // type spelled out, and the column type it reads as, where one does. struct FileColumn { @@ -42,7 +36,9 @@ bool ResolveColumnTypes(const std::vector& paths, Schema* schema, // Reads a parquet file into the record store one row group at a time, so the // Arrow buffers for a group are released before the next is read and the two -// representations are never both resident for the whole file. +// representations are never both resident for the whole file. The file arrives +// as a C Data stream from `parquet_io` and is decoded by `batch_loader`, so a +// file and a frame are read by the same code and nothing here links Arrow. bool LoadParquet(const std::string& path, const Schema& schema, RecordStore* store, LoadStats* stats, std::string* error); diff --git a/src/cpplink/predict.cpp b/src/cpplink/predict.cpp index b92634d..a70bfac 100644 --- a/src/cpplink/predict.cpp +++ b/src/cpplink/predict.cpp @@ -313,14 +313,28 @@ bool EdgeShardWriter::Close() { return !file_.fail(); } +void EdgeTable::Append(EdgeTable* other) { + a.insert(a.end(), other->a.begin(), other->a.end()); + b.insert(b.end(), other->b.begin(), other->b.end()); + gamma.insert(gamma.end(), other->gamma.begin(), other->gamma.end()); + weight.insert(weight.end(), other->weight.begin(), other->weight.end()); + *other = EdgeTable(); +} + bool Predict(const RecordStore& store, const ComparisonSet& comparisons, const BlockingPlan& plan, const Scorer& scorer, const PredictOptions& options, PredictReport* report, std::string* error) { const auto started = std::chrono::steady_clock::now(); const unsigned threads = ResolveThreads(options.threads); + const bool to_files = !options.out_dir.empty(); + if (!to_files && options.table == nullptr) { + *error = + "predict: nowhere to put the predictions, neither a directory nor a table"; + return false; + } std::error_code ec; - std::filesystem::create_directories(options.out_dir, ec); + if (to_files) std::filesystem::create_directories(options.out_dir, ec); if (ec) { *error = "cannot create output directory \"" + options.out_dir + "\": " + ec.message(); @@ -330,7 +344,7 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, const char* suffix = options.format == EdgeFormat::kBinary ? ".bin" : ".csv"; std::vector> writers; writers.reserve(threads); - for (unsigned t = 0; t < threads; ++t) { + for (unsigned t = 0; t < threads && to_files; ++t) { const std::string path = (std::filesystem::path(options.out_dir) / (EdgeShardName(t) + suffix)) .string(); @@ -385,6 +399,9 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, std::atomic emitted{0}; const uint64_t limit = options.max_edges; const bool csv = options.format == EdgeFormat::kCsv; + // One table per thread, joined after the threads do, for the same reason + // there is one shard file per thread. + std::vector tables(options.table != nullptr ? threads : 0); const std::vector sources = plan.AllSources(); std::unique_ptr progress; @@ -397,9 +414,10 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, plan, sources, threads, [&](unsigned t) { ThreadTally* counts = &tally[t]; - EdgeShardWriter* writer = writers[t].get(); + EdgeShardWriter* writer = to_files ? writers[t].get() : nullptr; SpillWriter* spill = spilling ? spills[t].get() : nullptr; - return [&, counts, writer, spill](uint32_t a, uint32_t b) { + EdgeTable* table = tables.empty() ? nullptr : &tables[t]; + return [&, counts, writer, spill, table](uint32_t a, uint32_t b) { Bump(&counts->enumerated); // Drawn per candidate, before anything decides the pair's fate, so // the sample stays uniform over candidates. @@ -452,6 +470,8 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, spill->Write(a, b, gamma); ++counts->spilled; } + if (table != nullptr) table->Push(a, b, gamma, weight); + if (writer == nullptr) return; if (csv) { writer->WriteCsv(store, a, b, gamma, weight); } else { @@ -463,7 +483,7 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, if (progress) progress->Finish(); for (unsigned t = 0; t < threads; ++t) { - if (!writers[t]->Close()) { + if (to_files && !writers[t]->Close()) { *error = "failed while writing \"" + report->shards[t] + "\""; return false; } @@ -482,6 +502,15 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, return false; } } + if (options.table != nullptr) { + uint64_t total = 0; + for (const EdgeTable& t : tables) total += t.Size(); + options.table->a.reserve(total); + options.table->b.reserve(total); + options.table->gamma.reserve(total); + options.table->weight.reserve(total); + for (EdgeTable& t : tables) options.table->Append(&t); + } report->threads = threads; report->mode = plan.mode(); report->datasets = plan.NumDatasets(); @@ -504,7 +533,7 @@ bool Predict(const RecordStore& store, const ComparisonSet& comparisons, manifest.layout = GammaLayout(comparisons); if (!WriteSpillManifest(options.spill_dir, manifest, error)) return false; } - if (!options.merge_path.empty()) { + if (!options.merge_path.empty() && to_files) { if (options.progress != nullptr) { *options.progress << " merging " << threads << (threads == 1 ? " shard" : " shards") << " into " @@ -567,7 +596,13 @@ void PrintPredictPlan(const RecordStore& store, const BlockingPlan& plan, } out << "\n"; const unsigned threads = ResolveThreads(options.threads); - if (!options.merge_path.empty()) { + const bool to_files = !options.out_dir.empty(); + if (!to_files) { + out << "Predictions kept in memory"; + if (options.table != nullptr) + out << " (one table per thread, joined at the end)"; + out << "\n"; + } else if (!options.merge_path.empty()) { out << "Predictions " << options.merge_path << " (" << threads << (threads == 1 ? " shard" : " shards") << " merged at the end)\n"; } else { @@ -682,6 +717,10 @@ void PrintPredictReport(const PredictReport& report, const Scorer& scorer, return; } out << "\nPredictions\n"; + if (report.shards.empty()) { + out << " kept in memory\n"; + return; + } for (const std::string& shard : report.shards) out << " " << shard << "\n"; } diff --git a/src/cpplink/predict.hpp b/src/cpplink/predict.hpp index 631ef65..1279b62 100644 --- a/src/cpplink/predict.hpp +++ b/src/cpplink/predict.hpp @@ -58,7 +58,28 @@ class EdgeShardWriter { std::string EdgeShardName(unsigned thread); +// Every prediction of a run held in memory, as parallel columns: what a front +// end hands out as a table with no file in between. Filled one shard per thread +// and joined at the end, so the rows are in the order the threads finished, as +// the shard files are. Twenty bytes a prediction, the shard's own width. +struct EdgeTable { + std::vector a; + std::vector b; + std::vector gamma; + std::vector weight; + + uint64_t Size() const { return a.size(); } + void Push(uint32_t row_a, uint32_t row_b, uint32_t pattern, double w) { + a.push_back(row_a); + b.push_back(row_b); + gamma.push_back(pattern); + weight.push_back(w); + } + void Append(EdgeTable* other); +}; + struct PredictOptions { + // The shard directory. Empty writes no shard at all, which needs `table`. std::string out_dir; // When set, `out_dir` is a staging directory: the per-thread shards are // written into it and merged into this one file at the end of the run, and @@ -77,6 +98,8 @@ struct PredictOptions { // re-tune upward and not enough to see what a lower threshold would find. double spill_sample = 0.0; uint64_t spill_seed = 20260904; + // Where to keep every prediction in memory, beside or instead of the shards. + EdgeTable* table = nullptr; // Where to report progress while the run walks the pair stream, or null for // none. With `progress_columns` set the stream is a terminal that wide and // one line is redrawn in place; without it a line is printed every few diff --git a/src/cpplink/rescore.cpp b/src/cpplink/rescore.cpp index 51f142f..b24ce1f 100644 --- a/src/cpplink/rescore.cpp +++ b/src/cpplink/rescore.cpp @@ -78,8 +78,14 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, return false; } + const bool to_files = !options.out_dir.empty(); + if (!to_files && options.table == nullptr) { + *error = + "rescore: nowhere to put the predictions, neither a directory nor a table"; + return false; + } std::error_code ec; - std::filesystem::create_directories(options.out_dir, ec); + if (to_files) std::filesystem::create_directories(options.out_dir, ec); if (ec) { *error = "cannot create output directory \"" + options.out_dir + "\": " + ec.message(); @@ -91,7 +97,7 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, const char* suffix = options.format == EdgeFormat::kBinary ? ".bin" : ".csv"; std::vector> writers; writers.reserve(threads); - for (unsigned t = 0; t < threads; ++t) { + for (unsigned t = 0; t < threads && to_files; ++t) { const std::string path = (std::filesystem::path(options.out_dir) / (EdgeShardName(t) + suffix)) .string(); @@ -112,10 +118,12 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, std::vector failures(threads); const uint64_t limit = options.max_edges; const bool csv = options.format == EdgeFormat::kCsv; + std::vector tables(options.table != nullptr ? threads : 0); auto work = [&](unsigned t) { ThreadTally* counts = &tally[t]; - EdgeShardWriter* writer = writers[t].get(); + EdgeShardWriter* writer = to_files ? writers[t].get() : nullptr; + EdgeTable* table = tables.empty() ? nullptr : &tables[t]; std::vector batch(kBatchPairs * 3); for (;;) { const size_t index = next.fetch_add(1); @@ -147,6 +155,8 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, } if (limit > 0 && emitted.fetch_add(1) >= limit) continue; ++counts->edges; + if (table != nullptr) table->Push(a, b, gamma, weight); + if (writer == nullptr) continue; if (csv) { writer->WriteCsv(store, a, b, gamma, weight); } else { @@ -174,7 +184,7 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, } } for (unsigned t = 0; t < threads; ++t) { - if (!writers[t]->Close()) { + if (to_files && !writers[t]->Close()) { *error = "failed while writing \"" + report->shards[t] + "\""; return false; } @@ -183,11 +193,14 @@ bool Rescore(const RecordStore& store, const ComparisonSet& comparisons, report->dropped += tally[t].dropped; report->tf_lookups += tally[t].tf_lookups; } + if (options.table != nullptr) { + for (EdgeTable& t : tables) options.table->Append(&t); + } report->threads = threads; report->truncated = limit > 0 && emitted.load() > limit; report->below_spill_threshold = report->manifest.sample_rate <= 0.0 && scorer.threshold() < report->manifest.threshold; - if (!options.merge_path.empty()) { + if (!options.merge_path.empty() && to_files) { if (!MergeStagedShards(store, options.out_dir, options.merge_path, options.format == EdgeFormat::kBinary, &report->merge_seconds, error)) { @@ -221,7 +234,9 @@ void PrintRescoreReport(const RescoreReport& report, const Scorer& scorer, << static_cast(report.pairs) / report.seconds << " pairs/s)"; } out << "\n"; - if (report.merged_path.empty()) { + if (report.merged_path.empty() && report.shards.empty()) { + out << " kept in memory\n"; + } else if (report.merged_path.empty()) { for (const std::string& shard : report.shards) out << " " << shard << "\n"; } else { out << " " << report.merged_path << " (" << report.threads << " shard" diff --git a/src/cpplink/rescore.hpp b/src/cpplink/rescore.hpp index 5b61777..6592b56 100644 --- a/src/cpplink/rescore.hpp +++ b/src/cpplink/rescore.hpp @@ -25,6 +25,9 @@ struct RescoreOptions { EdgeFormat format = EdgeFormat::kBinary; unsigned threads = 0; uint64_t max_edges = 0; + // Where to keep every prediction in memory, beside or instead of the shards, + // as `PredictOptions::table` does. An empty `out_dir` needs it. + EdgeTable* table = nullptr; }; struct RescoreReport { diff --git a/src/cpplink/sample_data.cpp b/src/cpplink/sample_data.cpp index 6dc54fa..5544424 100644 --- a/src/cpplink/sample_data.cpp +++ b/src/cpplink/sample_data.cpp @@ -12,9 +12,8 @@ #include #include -#include -#include -#include +#include "cpplink/arrow_export.hpp" +#include "cpplink/parquet_io.hpp" namespace cpplink { namespace { @@ -223,153 +222,89 @@ class Generator { Vocabulary street_words_; }; -std::shared_ptr MakeArrowSchema() { - return arrow::schema({ - arrow::field("id", arrow::utf8()), - arrow::field("first_name", arrow::utf8()), - arrow::field("last_name", arrow::utf8()), - arrow::field("gender", arrow::utf8()), - arrow::field("dob", arrow::date32()), - arrow::field("email", arrow::utf8()), - arrow::field("phone", arrow::utf8()), - arrow::field("postcode", arrow::utf8()), - arrow::field("latitude", arrow::float64()), - arrow::field("longitude", arrow::float64()), - arrow::field("address_tokens", arrow::list(arrow::utf8())), - }); -} - -// The eleven builders one output file accumulates rows into, and the writer they -// are flushed to. Routing a row to one of two files is then a choice of sink rather -// than a second copy of the append code. +// The eleven columns one output file accumulates rows into, as C Data batches, +// and the writer they are flushed to. Routing a row to one of two files is then +// a choice of sink rather than a second copy of the append code. struct RowSink { - explicit RowSink(arrow::MemoryPool* pool) - : id(pool), - first(pool), - last(pool), - gender(pool), - email(pool), - phone(pool), - postcode(pool), - dob(pool), - latitude(pool), - longitude(pool), - token_values(std::make_shared(pool)), - tokens(pool, token_values) {} - - bool Open(const std::string& path, const arrow::Schema& schema, std::string* error) { + RowSink() { + builder.AddColumn("id", ExportType::kString); + builder.AddColumn("first_name", ExportType::kString); + builder.AddColumn("last_name", ExportType::kString); + builder.AddColumn("gender", ExportType::kString); + builder.AddColumn("dob", ExportType::kDate32); + builder.AddColumn("email", ExportType::kString); + builder.AddColumn("phone", ExportType::kString); + builder.AddColumn("postcode", ExportType::kString); + builder.AddColumn("latitude", ExportType::kDouble); + builder.AddColumn("longitude", ExportType::kDouble); + builder.AddColumn("address_tokens", ExportType::kStringList); + } + + bool Open(const std::string& path, std::string* error) { this->path = path; - auto sink_result = arrow::io::FileOutputStream::Open(path); - if (!sink_result.ok()) { - *error = "cannot create " + path + ": " + sink_result.status().message(); - return false; - } - auto props = parquet::WriterProperties::Builder() - .compression(parquet::Compression::SNAPPY) - ->build(); - auto writer_result = parquet::arrow::FileWriter::Open( - schema, arrow::default_memory_pool(), *sink_result, props); - if (!writer_result.ok()) { - *error = "cannot open parquet writer for " + path + ": " + - writer_result.status().message(); - return false; - } - writer = std::move(*writer_result); - return true; + ArrowSchema schema; + builder.ExportSchema(&schema); + const bool ok = writer.Open(path, schema, error); + schema.release(&schema); + writing = ok; + return ok; } - bool Append(const SampleRecord& record, const std::string& identifier, - std::string* error) { - auto status = id.Append(identifier); - status &= first.Append(record.first_name); - status &= last.Append(record.last_name); - status &= record.gender < 0 ? gender.AppendNull() - : gender.Append(record.gender == 1 ? "M" : "F"); - status &= dob.Append(record.dob); - status &= record.email.empty() ? email.AppendNull() : email.Append(record.email); - status &= record.phone.empty() ? phone.AppendNull() : phone.Append(record.phone); - status &= record.postcode.empty() ? postcode.AppendNull() - : postcode.Append(record.postcode); - status &= latitude.Append(record.latitude); - status &= longitude.Append(record.longitude); - status &= tokens.Append(); - for (const std::string& token : record.address_tokens) { - status &= token_values->Append(token); + void Append(const SampleRecord& record, const std::string& identifier) { + builder.AppendString(0, identifier); + builder.AppendString(1, record.first_name); + builder.AppendString(2, record.last_name); + if (record.gender < 0) { + builder.AppendNull(3); + } else { + builder.AppendString(3, record.gender == 1 ? "M" : "F"); } - if (!status.ok()) { - *error = "building row " + identifier + ": " + status.message(); - return false; + builder.AppendDate32(4, record.dob); + if (record.email.empty()) { + builder.AppendNull(5); + } else { + builder.AppendString(5, record.email); } - ++pending; - return true; - } - - bool Flush(const std::shared_ptr& schema, std::string* error) { - if (pending == 0) return true; - std::vector> arrays(11); - arrow::Status status = id.Finish(&arrays[0]); - status &= first.Finish(&arrays[1]); - status &= last.Finish(&arrays[2]); - status &= gender.Finish(&arrays[3]); - status &= dob.Finish(&arrays[4]); - status &= email.Finish(&arrays[5]); - status &= phone.Finish(&arrays[6]); - status &= postcode.Finish(&arrays[7]); - status &= latitude.Finish(&arrays[8]); - status &= longitude.Finish(&arrays[9]); - status &= tokens.Finish(&arrays[10]); - if (!status.ok()) { - *error = "finishing a batch: " + status.message(); - return false; + if (record.phone.empty()) { + builder.AppendNull(6); + } else { + builder.AppendString(6, record.phone); } - const auto table = - arrow::Table::Make(schema, arrays, static_cast(pending)); - status = writer->WriteTable(*table, static_cast(pending)); - if (!status.ok()) { - *error = "writing a row group to " + path + ": " + status.message(); - return false; + if (record.postcode.empty()) { + builder.AppendNull(7); + } else { + builder.AppendString(7, record.postcode); } - pending = 0; - return true; + builder.AppendDouble(8, record.latitude); + builder.AppendDouble(9, record.longitude); + builder.AppendStringList(10, record.address_tokens); } - bool Close(std::string* error) { - const arrow::Status closed = writer->Close(); - if (!closed.ok()) { - *error = "closing " + path + ": " + closed.message(); - return false; - } - return true; + uint64_t Pending() const { return static_cast(builder.Rows()); } + + // A sink with no file keeps every row in its builder. + bool Flush(std::string* error) { + if (!writing || builder.Rows() == 0) return true; + ArrowArray batch; + if (!builder.ExportBatch(&batch, error)) return false; + return writer.Write(&batch, error); } - arrow::StringBuilder id, first, last, gender, email, phone, postcode; - arrow::Date32Builder dob; - arrow::DoubleBuilder latitude, longitude; - std::shared_ptr token_values; - arrow::ListBuilder tokens; - uint64_t pending = 0; + bool Close(std::string* error) { return !writing || writer.Close(error); } + + BatchBuilder builder; std::string path; - std::unique_ptr writer; + bool writing = false; + ParquetWriter writer; }; -} // namespace - -bool WriteSampleParquet(const std::string& path, const SampleOptions& options, - std::string* error) { - const auto schema = MakeArrowSchema(); - auto* pool = arrow::default_memory_pool(); - - // One sink per output file. With more than one file every planted duplicate is - // routed to one of the later ones, so the files are exactly the link fixture - // the cross-dataset path needs: every recorded pair crosses them. - std::vector> sinks; - sinks.push_back(std::make_unique(pool)); - if (!sinks.back()->Open(path, *schema, error)) return false; - const bool linking = !options.link_paths.empty(); - for (const std::string& link_path : options.link_paths) { - sinks.push_back(std::make_unique(pool)); - if (!sinks.back()->Open(link_path, *schema, error)) return false; - } +// Plants the rows into the sinks: originals into the first, every duplicate into +// one of the later ones where there are any. The one generator behind both the +// files and the in-memory tables. +bool GenerateInto(const SampleOptions& options, + std::vector>* sinks_ptr, std::string* error) { + std::vector>& sinks = *sinks_ptr; + const bool linking = sinks.size() > 1; // Never asked when there is no later file, but constructed either way, and a // distribution over an empty range is undefined. std::uniform_int_distribution which_link( @@ -429,18 +364,52 @@ bool WriteSampleParquet(const std::string& path, const SampleOptions& options, // A duplicate lands in one of the later files, drawn from the planner so // the file a row lands in is as reproducible as the row. RowSink& sink = *sinks[linking && duplicate ? which_link(planner) : 0]; - if (!sink.Append(record, identifier, error)) return false; - if (sink.pending >= static_cast(options.row_group_size) && - !sink.Flush(schema, error)) { + sink.Append(record, identifier); + if (sink.Pending() >= static_cast(options.row_group_size) && + !sink.Flush(error)) { return false; } } for (const std::unique_ptr& sink : sinks) { - if (!sink->Flush(schema, error)) return false; + if (!sink->Flush(error)) return false; + } + return true; +} + +} // namespace + +bool WriteSampleParquet(const std::string& path, const SampleOptions& options, + std::string* error) { + // One sink per output file. With more than one file every planted duplicate is + // routed to one of the later ones, so the files are exactly the link fixture + // the cross-dataset path needs: every recorded pair crosses them. + std::vector> sinks; + sinks.push_back(std::make_unique()); + if (!sinks.back()->Open(path, error)) return false; + for (const std::string& link_path : options.link_paths) { + sinks.push_back(std::make_unique()); + if (!sinks.back()->Open(link_path, error)) return false; + } + if (!GenerateInto(options, &sinks, error)) return false; + for (const std::unique_ptr& sink : sinks) { if (!sink->Close(error)) return false; } return true; } +bool GenerateSampleTables(const SampleOptions& options, std::vector* tables, + std::string* error) { + std::vector> sinks; + sinks.push_back(std::make_unique()); + for (size_t i = 0; i < options.link_paths.size(); ++i) { + sinks.push_back(std::make_unique()); + } + if (!GenerateInto(options, &sinks, error)) return false; + tables->clear(); + for (std::unique_ptr& sink : sinks) + tables->push_back(std::move(sink->builder)); + return true; +} + } // namespace cpplink diff --git a/src/cpplink/sample_data.hpp b/src/cpplink/sample_data.hpp index 4fbbb87..72b991e 100644 --- a/src/cpplink/sample_data.hpp +++ b/src/cpplink/sample_data.hpp @@ -7,6 +7,8 @@ #include #include +#include "cpplink/arrow_export.hpp" + namespace cpplink { struct SampleOptions { @@ -34,4 +36,11 @@ struct SampleOptions { bool WriteSampleParquet(const std::string& path, const SampleOptions& options, std::string* error); +// The sample as batches in memory rather than files: one `BatchBuilder` per +// output the options name, holding every row that output would have been +// written, and the truth file still written where `truth_path` says. What a +// front end returns as frames when no file is wanted. +bool GenerateSampleTables(const SampleOptions& options, std::vector* tables, + std::string* error); + } // namespace cpplink diff --git a/src/cpplink/waterfall.cpp b/src/cpplink/waterfall.cpp index dc7bde7..7e7162b 100644 --- a/src/cpplink/waterfall.cpp +++ b/src/cpplink/waterfall.cpp @@ -15,14 +15,12 @@ #include #include -#include -#include -#include -#include - +#include "cpplink/arrow_export.hpp" #include "cpplink/explain.hpp" #include "cpplink/format.hpp" #include "cpplink/id_index.hpp" +#include "cpplink/merge_edges.hpp" +#include "cpplink/parquet_io.hpp" namespace cpplink { namespace { @@ -72,107 +70,27 @@ bool ForEachCsvPrediction(const std::string& path, const PredictionSink& sink, return true; } +// Every row of the merged parquet, through the same reader clustering uses. A +// row with a null id is skipped, as the csv reader skips what it cannot parse. bool ForEachParquetPrediction(const std::string& path, const PredictionSink& sink, std::string* error) { - auto input = arrow::io::ReadableFile::Open(path); - if (!input.ok()) { - *error = "explain: cannot open " + path + ": " + input.status().message(); - return false; - } - parquet::arrow::FileReaderBuilder builder; - arrow::Status status = builder.Open(*input); - if (!status.ok()) { - *error = - "explain: cannot read parquet metadata for " + path + ": " + status.message(); - return false; - } - auto reader_result = builder.Build(); - if (!reader_result.ok()) { - *error = "explain: cannot open parquet reader for " + path + ": " + - reader_result.status().message(); - return false; - } - std::unique_ptr reader = std::move(*reader_result); - std::shared_ptr schema; - status = reader->GetSchema(&schema); - if (!status.ok()) { - *error = "explain: cannot read the schema of " + path + ": " + status.message(); - return false; - } - std::vector indices; - for (const char* name : {"id_a", "id_b", "gamma"}) { - const int at = schema->GetFieldIndex(name); - if (at < 0) { - *error = "explain: '" + path + "' has no column \"" + name + - "\", so it is not a cpplink prediction file"; - return false; - } - indices.push_back(at); - } - const bool datasets = schema->GetFieldIndex("dataset_a") >= 0; - if (datasets) { - for (const char* name : {"dataset_a", "dataset_b"}) { - const int at = schema->GetFieldIndex(name); - if (at < 0) { - *error = "explain: '" + path + "' has dataset_a but no " + name; - return false; - } - indices.push_back(at); - } - } - for (int group = 0; group < reader->num_row_groups(); ++group) { - auto group_result = reader->ReadRowGroup(group, indices); - if (!group_result.ok()) { - *error = "explain: cannot read row group " + std::to_string(group) + " of " + - path + ": " + group_result.status().message(); - return false; - } - const std::shared_ptr table = *group_result; - const auto ids_a = table->GetColumnByName("id_a"); - const auto ids_b = table->GetColumnByName("id_b"); - const auto gammas = table->GetColumnByName("gamma"); - const auto sets_a = datasets ? table->GetColumnByName("dataset_a") : nullptr; - const auto sets_b = datasets ? table->GetColumnByName("dataset_b") : nullptr; - for (int chunk = 0; chunk < ids_a->num_chunks(); ++chunk) { - const auto a_array = - std::dynamic_pointer_cast(ids_a->chunk(chunk)); - const auto b_array = - std::dynamic_pointer_cast(ids_b->chunk(chunk)); - const auto gamma_array = - std::dynamic_pointer_cast(gammas->chunk(chunk)); - std::shared_ptr set_a; - std::shared_ptr set_b; - if (datasets) { - set_a = - std::dynamic_pointer_cast(sets_a->chunk(chunk)); - set_b = - std::dynamic_pointer_cast(sets_b->chunk(chunk)); - } - if (a_array == nullptr || b_array == nullptr || gamma_array == nullptr || - (datasets && (set_a == nullptr || set_b == nullptr))) { - *error = "explain: '" + path + - "' holds id_a, id_b, dataset_a, dataset_b or gamma in a type " - "a cpplink prediction file does not use"; - return false; + EdgeColumns columns; + return ReadPredictionFile( + path, "explain", {"gamma"}, &columns, + [&](const EdgeRow& edge, std::string* visit_error) { + if (edge.id_a.empty() || edge.id_b.empty()) return true; + if (columns.datasets && (edge.dataset_a.empty() || edge.dataset_b.empty())) { + return true; } - for (int64_t row = 0; row < a_array->length(); ++row) { - if (a_array->IsNull(row) || b_array->IsNull(row) || - (datasets && (set_a->IsNull(row) || set_b->IsNull(row)))) { - continue; - } - NamedPrediction prediction; - if (datasets) { - prediction.dataset_a = set_a->GetView(row); - prediction.dataset_b = set_b->GetView(row); - } - prediction.id_a = a_array->GetView(row); - prediction.id_b = b_array->GetView(row); - prediction.gamma = gamma_array->IsNull(row) ? 0 : gamma_array->Value(row); - if (!sink(prediction, error)) return false; - } - } - } - return true; + NamedPrediction prediction; + prediction.dataset_a = edge.dataset_a; + prediction.id_a = edge.id_a; + prediction.dataset_b = edge.dataset_b; + prediction.id_b = edge.id_b; + prediction.gamma = edge.gamma; + return sink(prediction, visit_error); + }, + error); } // Column names are the comparison's name with a suffix, and an interaction's is @@ -259,97 +177,77 @@ class CsvWideSink : public WideSink { std::string buffer_; }; +// The wide row as C Data batches handed to the parquet writer: the same builder +// a data frame is handed in Python, so the two cannot name a column differently. class ParquetWideSink : public WideSink { public: ParquetWideSink(size_t batch_rows, size_t comparisons, size_t interactions) : batch_rows_(batch_rows == 0 ? 1 : batch_rows), - levels_(comparisons), - bits_(comparisons), - tf_(comparisons), - frequency_(comparisons), - interaction_(interactions) {} + comparisons_(comparisons), + interactions_(interactions) {} bool Open(const std::string& path, const std::vector& columns, std::string* error) { path_ = path; - std::vector> fields = { - arrow::field("id_a", arrow::utf8()), - arrow::field("id_b", arrow::utf8()), - arrow::field("gamma", arrow::uint32()), - arrow::field("prior", arrow::float64()), - arrow::field("match_weight", arrow::float64()), - arrow::field("match_probability", arrow::float64()), - arrow::field("bracket_low", arrow::float64()), - arrow::field("bracket_high", arrow::float64()), - arrow::field("threshold", arrow::float64()), - arrow::field("zone", arrow::utf8()), - arrow::field("emitted", arrow::boolean()), - }; + builder_.AddColumn("id_a", ExportType::kString); + builder_.AddColumn("id_b", ExportType::kString); + builder_.AddColumn("gamma", ExportType::kUInt32); + builder_.AddColumn("prior", ExportType::kDouble); + builder_.AddColumn("match_weight", ExportType::kDouble); + builder_.AddColumn("match_probability", ExportType::kDouble); + builder_.AddColumn("bracket_low", ExportType::kDouble); + builder_.AddColumn("bracket_high", ExportType::kDouble); + builder_.AddColumn("threshold", ExportType::kDouble); + builder_.AddColumn("zone", ExportType::kString); + builder_.AddColumn("emitted", ExportType::kBoolean); size_t at = kFixedColumns; - for (size_t c = 0; c < levels_.size(); ++c) { - fields.push_back(arrow::field(columns[at++], arrow::uint8())); - fields.push_back(arrow::field(columns[at++], arrow::float64())); - fields.push_back(arrow::field(columns[at++], arrow::float64())); - fields.push_back(arrow::field(columns[at++], arrow::uint32())); - } - for (size_t i = 0; i < interaction_.size(); ++i) { - fields.push_back(arrow::field(columns[at++], arrow::float64())); - } - schema_ = arrow::schema(fields); - auto sink = arrow::io::FileOutputStream::Open(path); - if (!sink.ok()) { - *error = "explain: cannot create " + path + ": " + sink.status().message(); - return false; + for (size_t c = 0; c < comparisons_; ++c) { + builder_.AddColumn(columns[at++], ExportType::kUInt8); + builder_.AddColumn(columns[at++], ExportType::kDouble); + builder_.AddColumn(columns[at++], ExportType::kDouble); + builder_.AddColumn(columns[at++], ExportType::kUInt32); } - auto props = parquet::WriterProperties::Builder() - .compression(parquet::Compression::SNAPPY) - ->build(); - auto writer = parquet::arrow::FileWriter::Open( - *schema_, arrow::default_memory_pool(), *sink, props); - if (!writer.ok()) { - *error = "explain: cannot open parquet writer for " + path + ": " + - writer.status().message(); - return false; + for (size_t i = 0; i < interactions_; ++i) { + builder_.AddColumn(columns[at++], ExportType::kDouble); } - writer_ = std::move(*writer); - return true; + ArrowSchema schema; + builder_.ExportSchema(&schema); + const bool ok = writer_.Open(path, schema, error); + schema.release(&schema); + if (!ok) *error = "explain: " + *error; + return ok; } bool Write(const PairWaterfall& w, std::string* error) override { - arrow::Status status = id_a_.Append(w.id_a); - status &= id_b_.Append(w.id_b); - status &= gamma_.Append(w.gamma); - status &= prior_.Append(w.prior); - status &= weight_.Append(w.weight); - status &= probability_.Append(w.probability); - status &= low_.Append(w.bracket_low); - status &= high_.Append(w.bracket_high); - status &= threshold_.Append(w.threshold); - status &= zone_.Append(ZoneName(w.zone)); - status &= emitted_.Append(w.emitted); + int at = 0; + builder_.AppendString(at++, w.id_a); + builder_.AppendString(at++, w.id_b); + builder_.AppendUInt32(at++, w.gamma); + builder_.AppendDouble(at++, w.prior); + builder_.AppendDouble(at++, w.weight); + builder_.AppendDouble(at++, w.probability); + builder_.AppendDouble(at++, w.bracket_low); + builder_.AppendDouble(at++, w.bracket_high); + builder_.AppendDouble(at++, w.threshold); + builder_.AppendString(at++, ZoneName(w.zone)); + builder_.AppendBoolean(at++, w.emitted); for (size_t c = 0; c < w.steps.size(); ++c) { - status &= levels_[c].Append(w.steps[c].level); - status &= bits_[c].Append(w.steps[c].bits); - status &= tf_[c].Append(w.steps[c].tf); - status &= frequency_[c].Append(w.steps[c].frequency); + builder_.AppendUInt8(at++, w.steps[c].level); + builder_.AppendDouble(at++, w.steps[c].bits); + builder_.AppendDouble(at++, w.steps[c].tf); + builder_.AppendUInt32(at++, w.steps[c].frequency); } for (size_t i = 0; i < w.interactions.size(); ++i) { - status &= interaction_[i].Append(w.interactions[i].bits); - } - if (!status.ok()) { - *error = "explain: building a row: " + status.message(); - return false; + builder_.AppendDouble(at++, w.interactions[i].bits); } - ++pending_; - if (pending_ >= batch_rows_) return Flush(error); + if (builder_.Rows() >= static_cast(batch_rows_)) return Flush(error); return true; } bool Close(std::string* error) override { if (!Flush(error)) return false; - const arrow::Status closed = writer_->Close(); - if (!closed.ok()) { - *error = "explain: closing " + path_ + ": " + closed.message(); + if (!writer_.Close(error)) { + *error = "explain: " + *error; return false; } return true; @@ -357,61 +255,21 @@ class ParquetWideSink : public WideSink { private: bool Flush(std::string* error) { - if (pending_ == 0) return true; - std::vector> arrays; - arrays.reserve(static_cast(schema_->num_fields())); - arrow::Status status; - const auto finish = [&](arrow::ArrayBuilder* builder) { - std::shared_ptr array; - status &= builder->Finish(&array); - arrays.push_back(std::move(array)); - }; - finish(&id_a_); - finish(&id_b_); - finish(&gamma_); - finish(&prior_); - finish(&weight_); - finish(&probability_); - finish(&low_); - finish(&high_); - finish(&threshold_); - finish(&zone_); - finish(&emitted_); - for (size_t c = 0; c < levels_.size(); ++c) { - finish(&levels_[c]); - finish(&bits_[c]); - finish(&tf_[c]); - finish(&frequency_[c]); - } - for (auto& builder : interaction_) finish(&builder); - if (!status.ok()) { - *error = "explain: finishing a batch: " + status.message(); - return false; - } - const auto table = - arrow::Table::Make(schema_, arrays, static_cast(pending_)); - status = writer_->WriteTable(*table, static_cast(pending_)); - if (!status.ok()) { - *error = "explain: writing a row group to " + path_ + ": " + status.message(); + if (builder_.Rows() == 0) return true; + ArrowArray batch; + if (!builder_.ExportBatch(&batch, error) || !writer_.Write(&batch, error)) { + *error = "explain: " + *error; return false; } - pending_ = 0; return true; } size_t batch_rows_; - std::shared_ptr schema_; - arrow::StringBuilder id_a_, id_b_, zone_; - arrow::UInt32Builder gamma_; - arrow::DoubleBuilder prior_, weight_, probability_, low_, high_, threshold_; - arrow::BooleanBuilder emitted_; - std::vector levels_; - std::vector bits_, tf_; - std::vector frequency_; - std::vector interaction_; - size_t pending_ = 0; + size_t comparisons_; + size_t interactions_; + BatchBuilder builder_; std::string path_; - std::unique_ptr writer_; + ParquetWriter writer_; }; } // namespace diff --git a/tests/arrow_export_test.cpp b/tests/arrow_export_test.cpp new file mode 100644 index 0000000..8d8d285 --- /dev/null +++ b/tests/arrow_export_test.cpp @@ -0,0 +1,277 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +// The producer side of the C Data seam. A batch built here is what the parquet +// writer takes and what a data frame is handed, so its layout is checked three +// ways: Arrow C++ imports and validates it, the consumer reads it back into a +// store, and a parquet round trip through `parquet_io` returns the same rows. + +#include "cpplink/arrow_export.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "cpplink/batch_loader.hpp" +#include "cpplink/parquet_io.hpp" +#include "cpplink/record_store.hpp" +#include "cpplink/schema.hpp" + +namespace { + +cpplink::Schema Parse(const std::string& json) { + cpplink::Schema schema; + std::string error; + EXPECT_TRUE(cpplink::ParseSchema(json, &schema, &error)) << error; + return schema; +} + +// Three rows over every export type, the middle one null wherever a null is +// allowed, and a list cell that is empty. +void FillRows(cpplink::BatchBuilder* builder) { + builder->AppendString(0, "1"); + builder->AppendString(0, "2"); + builder->AppendString(0, "3"); + builder->AppendString(1, "ann"); + builder->AppendNull(1); + builder->AppendString(1, "ann"); + builder->AppendBoolean(2, true); + builder->AppendNull(2); + builder->AppendBoolean(2, false); + builder->AppendUInt8(3, 7); + builder->AppendNull(3); + builder->AppendUInt8(3, 255); + builder->AppendUInt32(4, 70000); + builder->AppendNull(4); + builder->AppendUInt32(4, 1); + builder->AppendUInt64(5, 1ull << 40); + builder->AppendNull(5); + builder->AppendUInt64(5, 2); + builder->AppendDouble(6, 0.5); + builder->AppendNull(6); + builder->AppendDouble(6, -1.5); + builder->AppendDate32(7, 19000); + builder->AppendNull(7); + builder->AppendDate32(7, 19001); + builder->AppendStringList(8, {"y", "x", "y"}); + builder->AppendStringList(8, {}); + builder->AppendNull(8); +} + +cpplink::BatchBuilder EveryType() { + cpplink::BatchBuilder builder; + builder.AddColumn("id", cpplink::ExportType::kString); + builder.AddColumn("name", cpplink::ExportType::kString); + builder.AddColumn("flag", cpplink::ExportType::kBoolean); + builder.AddColumn("small", cpplink::ExportType::kUInt8); + builder.AddColumn("gamma", cpplink::ExportType::kUInt32); + builder.AddColumn("size", cpplink::ExportType::kUInt64); + builder.AddColumn("x", cpplink::ExportType::kDouble); + builder.AddColumn("dob", cpplink::ExportType::kDate32); + builder.AddColumn("tags", cpplink::ExportType::kStringList); + return builder; +} + +const char* kSchema = R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}, + {"name":"flag","type":"boolean"}, + {"name":"small","type":"string"}, + {"name":"gamma","type":"string"}, + {"name":"size","type":"string"}, + {"name":"x","type":"double"}, + {"name":"dob","type":"date"}, + {"name":"tags","type":"string_list"}]})"; + +void ExpectRows(const cpplink::RecordStore& store) { + ASSERT_EQ(store.NumRecords(), 3u); + EXPECT_EQ(store.ids().Get(1), "2"); + const auto& name = std::get(store.column(0)); + EXPECT_EQ(name.dict.Value(name.ids[0]), "ann"); + EXPECT_EQ(name.ids[1], cpplink::kNullId); + EXPECT_EQ(name.ids[0], name.ids[2]); + const auto& flag = std::get(store.column(1)); + EXPECT_EQ(flag.values, (std::vector{1, cpplink::kNullBoolean, 0})); + const auto& small = std::get(store.column(2)); + EXPECT_EQ(small.dict.Value(small.ids[2]), "255"); + EXPECT_EQ(small.ids[1], cpplink::kNullId); + const auto& gamma = std::get(store.column(3)); + EXPECT_EQ(gamma.dict.Value(gamma.ids[0]), "70000"); + const auto& size = std::get(store.column(4)); + EXPECT_EQ(size.dict.Value(size.ids[0]), "1099511627776"); + const auto& x = std::get(store.column(5)); + EXPECT_EQ(x.values[0], 0.5); + EXPECT_TRUE(std::isnan(x.values[1])); + EXPECT_EQ(x.values[2], -1.5); + const auto& dob = std::get(store.column(6)); + EXPECT_EQ(dob.values, (std::vector{19000, cpplink::kNullDate, 19001})); + const auto& tags = std::get(store.column(7)); + EXPECT_EQ(tags.offsets, (std::vector{0, 2, 2, 2})); + EXPECT_EQ(tags.dict.Size(), 2u); +} + +} // namespace + +TEST(ArrowExport, ArrowValidatesTheLayout) { + cpplink::BatchBuilder builder = EveryType(); + FillRows(&builder); + ArrowSchema schema; + ArrowArray array; + std::string error; + builder.ExportSchema(&schema); + ASSERT_TRUE(builder.ExportBatch(&array, &error)) << error; + EXPECT_EQ(builder.Rows(), 0) << "the rows moved out"; + + auto imported = arrow::ImportRecordBatch(&array, &schema); + ASSERT_TRUE(imported.ok()) << imported.status().message(); + const auto& batch = *imported; + ASSERT_TRUE(batch->ValidateFull().ok()) << batch->ValidateFull().message(); + EXPECT_EQ(batch->num_rows(), 3); + EXPECT_EQ(batch->schema()->field(0)->type()->ToString(), "string"); + EXPECT_EQ(batch->schema()->field(2)->type()->ToString(), "bool"); + EXPECT_EQ(batch->schema()->field(3)->type()->ToString(), "uint8"); + EXPECT_EQ(batch->schema()->field(4)->type()->ToString(), "uint32"); + EXPECT_EQ(batch->schema()->field(5)->type()->ToString(), "uint64"); + EXPECT_EQ(batch->schema()->field(6)->type()->ToString(), "double"); + EXPECT_EQ(batch->schema()->field(7)->type()->ToString(), "date32[day]"); + EXPECT_EQ(batch->schema()->field(8)->type()->ToString(), "list"); + EXPECT_EQ(batch->column(1)->null_count(), 1); + EXPECT_EQ(batch->column(0)->null_count(), 0); + EXPECT_EQ(batch->column(0)->data()->buffers[0], nullptr) + << "a column with no null carries no validity bitmap"; + EXPECT_EQ(batch->column(8)->null_count(), 1); + const auto& tags = static_cast(*batch->column(8)); + EXPECT_EQ(tags.value_length(0), 3); + EXPECT_EQ(tags.value_length(1), 0); + EXPECT_TRUE(tags.IsNull(2)); + // Releasing the imported batch released the export; the structs say so. + EXPECT_EQ(array.release, nullptr); + EXPECT_EQ(schema.release, nullptr); +} + +TEST(ArrowExport, TheConsumerReadsWhatTheProducerWrote) { + cpplink::BatchBuilder builder = EveryType(); + FillRows(&builder); + ArrowSchema schema; + ArrowArray array; + std::string error; + builder.ExportSchema(&schema); + ASSERT_TRUE(builder.ExportBatch(&array, &error)) << error; + + const cpplink::Schema spec = Parse(kSchema); + cpplink::RecordStore store(spec); + ASSERT_TRUE(cpplink::AppendRecordBatch(schema, array, spec, &store, &error)) << error; + array.release(&array); + schema.release(&schema); + store.set_num_records(3); + store.Finalize(); + ExpectRows(store); +} + +TEST(ArrowExport, RefusesRaggedColumns) { + cpplink::BatchBuilder builder; + builder.AddColumn("a", cpplink::ExportType::kString); + builder.AddColumn("b", cpplink::ExportType::kDouble); + builder.AppendString(0, "x"); + builder.AppendString(0, "y"); + builder.AppendDouble(1, 1.0); + ArrowArray array; + std::string error; + EXPECT_FALSE(builder.ExportBatch(&array, &error)); + EXPECT_NE(error.find("\"b\" has 1 rows where \"a\" has 2"), std::string::npos) + << error; + EXPECT_EQ(builder.Rows(), 2) << "a refused export leaves the rows in place"; +} + +TEST(ArrowExport, ABatchCanBeReusedAfterExport) { + cpplink::BatchBuilder builder; + builder.AddColumn("name", cpplink::ExportType::kString); + builder.AddColumn("tags", cpplink::ExportType::kStringList); + builder.AppendNull(0); + builder.AppendStringList(1, {"a"}); + ArrowArray first; + std::string error; + ASSERT_TRUE(builder.ExportBatch(&first, &error)) << error; + first.release(&first); + // The second batch starts from zero: fresh offsets, no validity bitmap. + builder.AppendString(0, "b"); + builder.AppendStringList(1, {"c", "d"}); + ArrowSchema schema; + ArrowArray second; + builder.ExportSchema(&schema); + ASSERT_TRUE(builder.ExportBatch(&second, &error)) << error; + auto imported = arrow::ImportRecordBatch(&second, &schema); + ASSERT_TRUE(imported.ok()) << imported.status().message(); + ASSERT_TRUE((*imported)->ValidateFull().ok()); + EXPECT_EQ((*imported)->column(0)->null_count(), 0); + const auto& names = static_cast(*(*imported)->column(0)); + EXPECT_EQ(names.GetString(0), "b"); + const auto& tags = static_cast(*(*imported)->column(1)); + EXPECT_EQ(tags.value_offset(0), 0); + EXPECT_EQ(tags.value_length(0), 2); +} + +// The io seam end to end: batches written as row groups and read back as the +// same batches, through nothing but the C structs on either side. +TEST(ArrowExport, ParquetRoundTrip) { + const auto dir = std::filesystem::temp_directory_path() / + ("cpplink_arrow_export_" + std::to_string(::getpid())); + std::filesystem::create_directories(dir); + const std::string path = (dir / "rows.parquet").string(); + + ASSERT_TRUE(cpplink::ParquetSupported()); + cpplink::BatchBuilder builder = EveryType(); + cpplink::ParquetWriter writer; + ArrowSchema schema; + std::string error; + builder.ExportSchema(&schema); + ASSERT_TRUE(writer.Open(path, schema, &error)) << error; + schema.release(&schema); + for (int batch = 0; batch < 2; ++batch) { + FillRows(&builder); + ArrowArray array; + ASSERT_TRUE(builder.ExportBatch(&array, &error)) << error; + ASSERT_TRUE(writer.Write(&array, &error)) << error; + } + ASSERT_TRUE(writer.Close(&error)) << error; + + // Read only some columns, in an order of our choosing. + ArrowSchema file; + ASSERT_TRUE(cpplink::ReadParquetSchema(path, &file, &error)) << error; + EXPECT_EQ(file.n_children, 9); + EXPECT_EQ(cpplink::ArrowTypeName(*file.children[8]), "list") + << "parquet names the child element, and the name follows the file"; + file.release(&file); + + ArrowArrayStream stream; + ASSERT_TRUE(cpplink::OpenParquetStream(path, {"tags", "id", "x"}, &stream, &error)) + << error; + const cpplink::Schema spec = Parse(R"({"unique_id":"id","columns":[ + {"name":"x","type":"double"}, + {"name":"tags","type":"string_list"}]})"); + cpplink::RecordStore store(spec); + uint64_t rows = 0; + int batches = 0; + ASSERT_TRUE( + cpplink::AppendArrayStream(&stream, spec, &store, &rows, &batches, &error)) + << error; + EXPECT_EQ(rows, 6u); + EXPECT_EQ(batches, 2) << "one row group per batch written"; + store.set_num_records(6); + store.Finalize(); + EXPECT_EQ(store.ids().Get(5), "3"); + const auto& tags = std::get(store.column(1)); + EXPECT_EQ(tags.offsets, (std::vector{0, 2, 2, 2, 4, 4, 4})); + + ArrowArrayStream missing; + EXPECT_FALSE(cpplink::OpenParquetStream(path, {"gone"}, &missing, &error)); + EXPECT_NE(error.find("\"gone\" is not in"), std::string::npos) << error; + std::filesystem::remove_all(dir); +} diff --git a/tests/batch_loader_test.cpp b/tests/batch_loader_test.cpp new file mode 100644 index 0000000..931654a --- /dev/null +++ b/tests/batch_loader_test.cpp @@ -0,0 +1,536 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +// The batch loader reads Arrow buffers where a producer left them, through the C +// Data Interface, and it is the one decoder both the parquet path and a data frame +// go through. The parquet tests already cover the types a file holds; these cover +// what only a frame produces -- a sliced array, a sliced batch, a category +// column, string views, an all-null column, a stream of batches -- and check that +// a table decoded from its buffers is the store the parquet path builds from the +// same table, member for member. + +#include "cpplink/batch_loader.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "cpplink/parquet_loader.hpp" +#include "cpplink/record_store.hpp" +#include "cpplink/schema.hpp" + +namespace { + +template +std::shared_ptr Column(const std::vector& values, + const std::vector& valid = {}) { + Builder builder; + for (size_t i = 0; i < values.size(); ++i) { + if (!valid.empty() && !valid[i]) { + EXPECT_TRUE(builder.AppendNull().ok()); + } else { + EXPECT_TRUE(builder.Append(values[i]).ok()); + } + } + std::shared_ptr array; + EXPECT_TRUE(builder.Finish(&array).ok()); + return array; +} + +std::string Text(const cpplink::StringColumn& column, uint64_t row) { + return std::string(column.dict.Value(column.ids[row])); +} + +cpplink::Schema Parse(const std::string& json) { + cpplink::Schema schema; + std::string error; + EXPECT_TRUE(cpplink::ParseSchema(json, &schema, &error)) << error; + return schema; +} + +// Exports the batch and decodes it, releasing the structs afterwards as a +// consumer must. +bool Append(const arrow::RecordBatch& batch, const cpplink::Schema& schema, + cpplink::RecordStore* store, std::string* error) { + ArrowArray array; + ArrowSchema c_schema; + EXPECT_TRUE(arrow::ExportRecordBatch(batch, &array, &c_schema).ok()); + const bool ok = cpplink::AppendRecordBatch(c_schema, array, schema, store, error); + array.release(&array); + c_schema.release(&c_schema); + return ok; +} + +std::shared_ptr Batch( + const std::vector>& fields, + const std::vector>& arrays) { + return arrow::RecordBatch::Make(arrow::schema(fields), arrays.front()->length(), + arrays); +} + +// Every member of every column, so two stores built two ways can be compared. +void ExpectSameStore(const cpplink::RecordStore& a, const cpplink::RecordStore& b) { + ASSERT_EQ(a.NumRecords(), b.NumRecords()); + ASSERT_EQ(a.NumColumns(), b.NumColumns()); + EXPECT_EQ(a.ids().text, b.ids().text); + EXPECT_EQ(a.ids().offsets, b.ids().offsets); + for (size_t c = 0; c < a.NumColumns(); ++c) { + ASSERT_EQ(a.column(c).index(), b.column(c).index()); + std::visit( + [&](const auto& x) { + using T = std::decay_t; + const auto& y = std::get(b.column(c)); + if constexpr (std::is_same_v) { + EXPECT_EQ(x.ids, y.ids); + EXPECT_EQ(x.tf, y.tf); + ASSERT_EQ(x.dict.Size(), y.dict.Size()); + for (uint32_t v = 0; v < x.dict.Size(); ++v) { + EXPECT_EQ(x.dict.Value(v), y.dict.Value(v)); + } + } else if constexpr (std::is_same_v) { + EXPECT_EQ(x.ids, y.ids); + EXPECT_EQ(x.offsets, y.offsets); + EXPECT_EQ(x.tf, y.tf); + ASSERT_EQ(x.dict.Size(), y.dict.Size()); + for (uint32_t v = 0; v < x.dict.Size(); ++v) { + EXPECT_EQ(x.dict.Value(v), y.dict.Value(v)); + } + } else if constexpr (std::is_same_v) { + EXPECT_EQ(x.values, y.values); + EXPECT_EQ(x.tf, y.tf); + EXPECT_EQ(x.tf_origin, y.tf_origin); + } else if constexpr (std::is_same_v) { + EXPECT_EQ(x.values, y.values); + EXPECT_EQ(x.tf, y.tf); + } else { + ASSERT_EQ(x.values.size(), y.values.size()); + for (size_t i = 0; i < x.values.size(); ++i) { + if (std::isnan(x.values[i])) { + EXPECT_TRUE(std::isnan(y.values[i])); + } else { + EXPECT_EQ(x.values[i], y.values[i]); + } + } + } + }, + a.column(c)); + } +} + +} // namespace + +// A slice of a frame is the common way a wrapper bounds the transient copy of +// an object column: it exports arrays whose `offset` is not zero, and the values +// before the slice must not be read, let alone interned. +TEST(BatchLoader, ReadsSlicedArraysAtTheirOffset) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}, + {"name":"tags","type":"string_list"}, + {"name":"dob","type":"date"}, + {"name":"x","type":"double"}, + {"name":"flag","type":"boolean"}]})"); + + arrow::ListBuilder tags(arrow::default_memory_pool(), + std::make_shared()); + auto* tag = static_cast(tags.value_builder()); + for (const char* row : {"before", "b", "c", "after"}) { + ASSERT_TRUE(tags.Append().ok()); + ASSERT_TRUE(tag->AppendValues({row, "shared"}).ok()); + } + std::shared_ptr tags_array; + ASSERT_TRUE(tags.Finish(&tags_array).ok()); + + const auto whole = + Batch({arrow::field("id", arrow::utf8()), arrow::field("name", arrow::utf8()), + arrow::field("tags", arrow::list(arrow::utf8())), + arrow::field("dob", arrow::date32()), arrow::field("x", arrow::float64()), + arrow::field("flag", arrow::boolean())}, + {Column(std::vector{"0", "1", "2", "3"}), + Column( + std::vector{"zero", "one", "", "three"}, + {true, true, false, true}), + tags_array, Column(std::vector{1, 2, 3, 4}), + Column(std::vector{0.5, 1.5, 2.5, 3.5}), + Column(std::vector{false, true, false, true}, + {true, true, false, true})}); + + // Each column sliced on its own, so every array carries its own offset. + std::vector> sliced; + for (int c = 0; c < whole->num_columns(); ++c) { + sliced.push_back(whole->column(c)->Slice(1, 2)); + } + const auto batch = arrow::RecordBatch::Make(whole->schema(), 2, sliced); + + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + store.set_num_records(2); + store.Finalize(); + + EXPECT_EQ(store.ids().Get(0), "1"); + EXPECT_EQ(store.ids().Get(1), "2"); + const auto& name = std::get(store.column(0)); + EXPECT_EQ(Text(name, 0), "one"); + EXPECT_EQ(name.ids[1], cpplink::kNullId); + EXPECT_EQ(name.dict.Size(), 1u) << "values outside the slice are not interned"; + + const auto& tag_column = std::get(store.column(1)); + const std::vector offsets = {0, 2, 4}; + EXPECT_EQ(tag_column.offsets, offsets); + EXPECT_EQ(tag_column.dict.Size(), 3u) << "b, shared and c; not before or after"; + EXPECT_EQ(tag_column.tf[tag_column.ids[1]], 2u) << "shared is on both rows"; + + const auto& dob = std::get(store.column(2)); + EXPECT_EQ(dob.values, (std::vector{2, 3})); + const auto& x = std::get(store.column(3)); + EXPECT_EQ(x.values, (std::vector{1.5, 2.5})); + const auto& flag = std::get(store.column(4)); + EXPECT_EQ(flag.values, (std::vector{1, cpplink::kNullBoolean})); +} + +// A batch sliced as a whole is a struct with an offset, which applies to every +// child: the children keep their full length and the batch says where to read. +TEST(BatchLoader, AppliesTheBatchOffsetToEveryColumn) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}]})"); + const auto whole = Batch( + {arrow::field("id", arrow::int64()), arrow::field("name", arrow::utf8())}, + {Column(std::vector{10, 11, 12, 13}), + Column(std::vector{"a", "b", "c", "d"})}); + const auto batch = whole->Slice(2, 2); + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + ASSERT_EQ(store.ids().offsets.size(), 3u); + EXPECT_EQ(store.ids().Get(0), "12"); + EXPECT_EQ(store.ids().Get(1), "13"); + const auto& name = std::get(store.column(0)); + EXPECT_EQ(Text(name, 0), "c"); + EXPECT_EQ(Text(name, 1), "d"); +} + +// A pandas `category` arrives dictionary-encoded, which is the store's own shape: +// each distinct value is interned once and the rows are a remap of the indices. +// A category no row uses is not interned, or it would count as a distinct value +// and sit in every dictionary self-join with a frequency of zero. +TEST(BatchLoader, DictionaryColumnsInternEachValueOnce) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"city","type":"string"}, + {"name":"code","type":"string"}, + {"name":"tags","type":"string_list"}]})"); + + const auto cities = + Column(std::vector{"paris", "rome", "oslo"}); + const auto city_indices = Column(std::vector{1, 0, 1, 0}, + {true, true, false, true}); + const auto city = std::make_shared( + arrow::dictionary(arrow::int8(), arrow::utf8()), city_indices, cities); + + // An integer dictionary spells its digits once. + const auto codes = Column(std::vector{2000, 90210}); + const auto code_indices = + Column(std::vector{1, 1, 0, 0}); + const auto code = std::make_shared( + arrow::dictionary(arrow::int32(), arrow::int32()), code_indices, codes); + + // Dictionary-encoded list elements go through the same text visitor. + const auto elements = + Column(std::vector{"x", "y"}); + const auto element_indices = + Column(std::vector{1, 0, 0, 1, 1}); + const auto element = std::make_shared( + arrow::dictionary(arrow::int16(), arrow::utf8()), element_indices, elements); + const auto tags = std::make_shared( + arrow::list(element->type()), 4, + Column(std::vector{0, 2, 2, 3, 5}) + ->data() + ->buffers[1], + element); + + const auto batch = + Batch({arrow::field("id", arrow::utf8()), arrow::field("city", city->type()), + arrow::field("code", code->type()), arrow::field("tags", tags->type())}, + {Column(std::vector{"a", "b", "c", "d"}), + city, code, tags}); + + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*batch->schema(), &c_schema).ok()); + cpplink::ColumnType type; + EXPECT_TRUE(cpplink::ColumnTypeOf(*c_schema.children[1], &type)); + EXPECT_EQ(type, cpplink::ColumnType::kString); + EXPECT_TRUE(cpplink::ColumnTypeOf(*c_schema.children[2], &type)); + EXPECT_EQ(type, cpplink::ColumnType::kString) << "an integer dictionary is text"; + EXPECT_EQ(cpplink::ArrowTypeName(*c_schema.children[1]), + "dictionary"); + c_schema.release(&c_schema); + + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + store.set_num_records(4); + store.Finalize(); + + const auto& city_column = std::get(store.column(0)); + EXPECT_EQ(city_column.dict.Size(), 2u) << "oslo is a category no row holds"; + EXPECT_EQ(Text(city_column, 0), "rome"); + EXPECT_EQ(Text(city_column, 1), "paris"); + EXPECT_EQ(city_column.ids[2], cpplink::kNullId); + EXPECT_EQ(city_column.ids[1], city_column.ids[3]); + EXPECT_EQ(city_column.tf[city_column.ids[0]], 1u); + + const auto& code_column = std::get(store.column(1)); + EXPECT_EQ(Text(code_column, 0), "90210"); + EXPECT_EQ(Text(code_column, 3), "2000"); + EXPECT_EQ(code_column.dict.Size(), 2u); + + const auto& tag_column = std::get(store.column(2)); + const std::vector offsets = {0, 2, 2, 3, 4}; + EXPECT_EQ(tag_column.offsets, offsets) << "the last cell's y, y is one y"; + EXPECT_EQ(tag_column.dict.Value(tag_column.ids[3]), "y"); +} + +// polars hands text over as string views: short values inline, long ones in a +// data buffer the view points into. +TEST(BatchLoader, ReadsStringViews) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}]})"); + const std::string long_value(40, 'q'); + const auto batch = + Batch({arrow::field("id", arrow::utf8_view()), + arrow::field("name", arrow::utf8_view())}, + {Column(std::vector{"1", "2", "3"}), + Column( + std::vector{"short", long_value, "short"}, + {true, true, true})}); + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*batch->schema(), &c_schema).ok()); + EXPECT_EQ(cpplink::ArrowTypeName(*c_schema.children[1]), "string_view"); + cpplink::ColumnType type; + EXPECT_TRUE(cpplink::ColumnTypeOf(*c_schema.children[1], &type)); + EXPECT_EQ(type, cpplink::ColumnType::kString); + c_schema.release(&c_schema); + + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + EXPECT_EQ(store.ids().Get(2), "3"); + const auto& name = std::get(store.column(0)); + EXPECT_EQ(Text(name, 0), "short"); + EXPECT_EQ(Text(name, 1), long_value); + EXPECT_EQ(name.ids[0], name.ids[2]); +} + +// A column of nothing but None exports as the null type. It is missing on every +// row for whatever type the schema declares, not an error. +TEST(BatchLoader, NullTypedColumnsAreMissingEverywhere) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}, + {"name":"tags","type":"string_list"}, + {"name":"dob","type":"date"}, + {"name":"x","type":"double"}, + {"name":"flag","type":"boolean"}]})"); + const auto nulls = std::make_shared(2); + const auto batch = + Batch({arrow::field("id", arrow::utf8()), arrow::field("name", arrow::null()), + arrow::field("tags", arrow::null()), arrow::field("dob", arrow::null()), + arrow::field("x", arrow::null()), arrow::field("flag", arrow::null())}, + {Column(std::vector{"1", "2"}), nulls, + nulls, nulls, nulls, nulls}); + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + store.set_num_records(2); + store.Finalize(); + for (size_t c = 0; c < 5; ++c) EXPECT_EQ(store.NullCount(c), 2u) << c; + + // The type is not one a schema can take from the input, though. + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*batch->schema(), &c_schema).ok()); + cpplink::ColumnType type; + EXPECT_FALSE(cpplink::ColumnTypeOf(*c_schema.children[1], &type)); + EXPECT_EQ(cpplink::ArrowTypeName(*c_schema.children[1]), "null"); + c_schema.release(&c_schema); +} + +// Timestamps of every unit are days, as the parquet path has always read them, +// and the type is named the way Arrow prints it. +TEST(BatchLoader, TimestampsAreDays) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"seen","type":"date"}]})"); + arrow::TimestampBuilder builder(arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + arrow::default_memory_pool()); + ASSERT_TRUE(builder.Append(3 * 86400000000LL + 5).ok()); + ASSERT_TRUE(builder.AppendNull().ok()); + std::shared_ptr seen; + ASSERT_TRUE(builder.Finish(&seen).ok()); + const auto batch = + Batch({arrow::field("id", arrow::utf8()), arrow::field("seen", seen->type())}, + {Column(std::vector{"1", "2"}), seen}); + ArrowSchema c_schema; + ASSERT_TRUE(arrow::ExportSchema(*batch->schema(), &c_schema).ok()); + EXPECT_EQ(cpplink::ArrowTypeName(*c_schema.children[1]), "timestamp[us, tz=UTC]"); + EXPECT_EQ(cpplink::ArrowTypeName(*c_schema.children[1]), + batch->schema()->field(1)->type()->ToString()); + c_schema.release(&c_schema); + + cpplink::RecordStore store(schema); + std::string error; + ASSERT_TRUE(Append(*batch, schema, &store, &error)) << error; + const auto& column = std::get(store.column(0)); + EXPECT_EQ(column.values, (std::vector{3, cpplink::kNullDate})); +} + +TEST(BatchLoader, RefusesWhatTheSchemaCannotRead) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"x","type":"string"}]})"); + const auto batch = + Batch({arrow::field("id", arrow::utf8()), arrow::field("x", arrow::float64())}, + {Column(std::vector{"1"}), + Column(std::vector{1.0})}); + cpplink::RecordStore store(schema); + std::string error; + EXPECT_FALSE(Append(*batch, schema, &store, &error)); + EXPECT_NE(error.find("x: expected a string or integer column, found double"), + std::string::npos) + << error; + + const auto missing = Parse(R"({"unique_id":"id","columns":[ + {"name":"gone","type":"string"}]})"); + cpplink::RecordStore other(missing); + EXPECT_FALSE(Append(*batch, missing, &other, &error)); + EXPECT_NE(error.find("\"gone\" is not in the input"), std::string::npos) << error; +} + +// A stream is drained batch by batch, each released before the next is fetched. +TEST(BatchLoader, DrainsAStream) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}]})"); + const auto fields = arrow::schema( + {arrow::field("id", arrow::int64()), arrow::field("name", arrow::large_utf8())}); + std::vector> batches = { + arrow::RecordBatch::Make( + fields, 2, + {Column(std::vector{1, 2}), + Column(std::vector{"a", "b"})}), + arrow::RecordBatch::Make( + fields, 1, + {Column(std::vector{3}), + Column(std::vector{"a"})})}; + auto reader = arrow::RecordBatchReader::Make(batches, fields); + ASSERT_TRUE(reader.ok()); + ArrowArrayStream stream; + ASSERT_TRUE(arrow::ExportRecordBatchReader(*reader, &stream).ok()); + + cpplink::RecordStore store(schema); + uint64_t rows = 0; + int count = 0; + std::string error; + ASSERT_TRUE( + cpplink::AppendArrayStream(&stream, schema, &store, &rows, &count, &error)) + << error; + EXPECT_EQ(stream.release, nullptr) << "the stream is released on return"; + EXPECT_EQ(rows, 3u); + EXPECT_EQ(count, 2); + store.set_num_records(3); + store.Finalize(); + EXPECT_EQ(store.ids().Get(2), "3"); + const auto& name = std::get(store.column(0)); + EXPECT_EQ(name.ids[0], name.ids[2]); + EXPECT_EQ(name.tf[name.ids[0]], 2u); +} + +// The equivalence the design rests on: a table read from its buffers is the +// store the parquet path builds from the same table written to disk, member for +// member, across every column type. +TEST(BatchLoader, MatchesTheParquetPathMemberForMember) { + const auto schema = Parse(R"({"unique_id":"id","columns":[ + {"name":"name","type":"string"}, + {"name":"postcode","type":"string"}, + {"name":"tags","type":"string_list"}, + {"name":"dob","type":"date"}, + {"name":"x","type":"double"}, + {"name":"flag","type":"boolean"}]})"); + + arrow::LargeListBuilder tags(arrow::default_memory_pool(), + std::make_shared()); + auto* tag = static_cast(tags.value_builder()); + ASSERT_TRUE(tags.Append().ok()); + ASSERT_TRUE(tag->AppendValues({"y", "x", "y"}).ok()); + ASSERT_TRUE(tags.AppendNull().ok()); + ASSERT_TRUE(tags.Append().ok()); + ASSERT_TRUE(tags.Append().ok()); + ASSERT_TRUE(tag->AppendValues({"x"}).ok()); + std::shared_ptr tags_array; + ASSERT_TRUE(tags.Finish(&tags_array).ok()); + + arrow::TimestampBuilder dob_builder(arrow::timestamp(arrow::TimeUnit::NANO), + arrow::default_memory_pool()); + ASSERT_TRUE(dob_builder + .AppendValues({86400000000000LL * 10, 0, 86400000000000LL * 10, + 86400000000000LL * 11}, + std::vector{true, false, true, true}) + .ok()); + std::shared_ptr dob; + ASSERT_TRUE(dob_builder.Finish(&dob).ok()); + + const auto table = arrow::Table::Make( + arrow::schema({arrow::field("id", arrow::int64()), + arrow::field("name", arrow::large_utf8()), + arrow::field("postcode", arrow::int32()), + arrow::field("tags", arrow::large_list(arrow::large_utf8())), + arrow::field("dob", dob->type()), + arrow::field("x", arrow::float32()), + arrow::field("flag", arrow::boolean())}), + {Column(std::vector{1, 2, 3, 4}), + Column( + std::vector{"ann", "bob", "", "ann"}, + {true, true, false, true}), + Column(std::vector{2000, 0, 2000, 3}, + {true, false, true, true}), + tags_array, dob, + Column(std::vector{0.5f, 0.f, 1.5f, 2.5f}, + {true, false, true, true}), + Column(std::vector{true, false, false, true}, + {true, true, false, true})}); + + const auto dir = std::filesystem::temp_directory_path() / + ("cpplink_batch_loader_" + std::to_string(::getpid())); + std::filesystem::create_directories(dir); + const std::string path = (dir / "table.parquet").string(); + auto sink = arrow::io::FileOutputStream::Open(path); + ASSERT_TRUE(sink.ok()); + ASSERT_TRUE(parquet::arrow::WriteTable(*table, arrow::default_memory_pool(), *sink, + /*chunk_size=*/3) + .ok()); + + cpplink::RecordStore from_file(schema); + std::string error; + ASSERT_TRUE(cpplink::LoadParquet(path, schema, &from_file, nullptr, &error)) << error; + std::filesystem::remove_all(dir); + + // The same table, decoded from its buffers in batches of two. + cpplink::RecordStore from_buffers(schema); + arrow::TableBatchReader reader(*table); + reader.set_chunksize(2); + std::shared_ptr batch; + while (reader.ReadNext(&batch).ok() && batch != nullptr) { + ASSERT_TRUE(Append(*batch, schema, &from_buffers, &error)) << error; + } + from_buffers.set_num_records(4); + from_buffers.Finalize(); + + ExpectSameStore(from_file, from_buffers); + EXPECT_EQ(from_buffers.ids().Get(3), "4"); + const auto& dob_column = std::get(from_buffers.column(3)); + EXPECT_EQ(dob_column.values[0], 10); +} diff --git a/tests/edge_table_test.cpp b/tests/edge_table_test.cpp new file mode 100644 index 0000000..9ab2bdc --- /dev/null +++ b/tests/edge_table_test.cpp @@ -0,0 +1,354 @@ +// Copyright 2026 Mathieu Fourment +// SPDX-License-Identifier: MIT + +// The in-memory side of a run: predictions kept as a table rather than written +// as shards, clustered from that table or from a prediction table that arrived +// as a C Data stream, and the table exported with the record ids as a +// dictionary over the store's own id arena. Each is held to what the file path +// does, because a front end that returns a frame must return the file's rows. + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "cpplink/arrow_export.hpp" +#include "cpplink/blocking.hpp" +#include "cpplink/cluster.hpp" +#include "cpplink/comparison.hpp" +#include "cpplink/model.hpp" +#include "cpplink/predict.hpp" +#include "cpplink/record_store.hpp" +#include "cpplink/schema.hpp" +#include "cpplink/score.hpp" + +namespace { + +constexpr uint64_t kRecords = 400; + +const char* const kSchemaJson = R"({ + "unique_id": "id", + "columns": [ + {"name": "surname", "type": "string"}, + {"name": "city", "type": "string"} + ], + "comparisons": [ + {"name": "surname", "columns": ["surname"], "term_frequency": true, + "levels": [{"type": "null"}, {"type": "exact"}, {"type": "else"}]}, + {"name": "city", "columns": ["city"], + "levels": [{"type": "exact"}, {"type": "else"}]} + ], + "blocking": [{"type": "exact_value", "column": "city"}] +})"; + +struct Edge { + uint32_t a = 0; + uint32_t b = 0; + uint32_t gamma = 0; + double weight = 0.0; + bool operator<(const Edge& other) const { + return std::tie(a, b, gamma) < std::tie(other.a, other.b, other.gamma); + } + bool operator==(const Edge& other) const { + return a == other.a && b == other.b && gamma == other.gamma && + weight == other.weight; + } +}; + +class EdgeTableFixture : public ::testing::Test { + protected: + void SetUp() override { + dir_ = std::filesystem::temp_directory_path() / + ("cpplink_edge_table_" + std::to_string(::getpid())); + std::filesystem::remove_all(dir_); + std::filesystem::create_directories(dir_); + std::string error; + ASSERT_TRUE(cpplink::ParseSchema(kSchemaJson, &schema_, &error)) << error; + store_ = std::make_unique(schema_); + auto& surname = std::get(store_->mutable_column(0)); + auto& city = std::get(store_->mutable_column(1)); + std::vector names; + for (int i = 0; i < 40; ++i) { + names.push_back(surname.dict.Intern("name" + std::to_string(i))); + } + std::vector cities; + for (int i = 0; i < 4; ++i) { + cities.push_back(city.dict.Intern("city" + std::to_string(i))); + } + for (uint64_t row = 0; row < kRecords; ++row) { + const size_t pick = static_cast((row * row) % 40) / 2; + surname.ids.push_back(names[pick]); + city.ids.push_back(cities[(row * 7) % 4]); + store_->mutable_ids().Append("r" + std::to_string(row)); + } + store_->set_num_records(kRecords); + store_->Finalize(); + ASSERT_TRUE(plan_.Build(schema_, *store_, &error)) << error; + ASSERT_TRUE(comparisons_.Bind(schema_, *store_, &error)) << error; + model_.lambda = 0.05; + model_.records = kRecords; + model_.comparisons = { + Comparison("surname", true, {{1e-9, 1e-9}, {0.9, 0.02}, {0.1, 0.98}}), + Comparison("city", false, {{0.8, 0.25}, {0.2, 0.75}})}; + cpplink::ScoreOptions score; + score.threshold = -2.0; + ASSERT_TRUE(scorer_.Bind(model_, comparisons_, *store_, score, &error)) << error; + } + + void TearDown() override { + std::error_code ec; + std::filesystem::remove_all(dir_, ec); + } + + static cpplink::ModelComparison Comparison( + const std::string& name, bool tf, + const std::vector>& levels) { + cpplink::ModelComparison comparison; + comparison.name = name; + comparison.term_frequency = tf; + for (const auto& entry : levels) { + cpplink::ModelLevel level; + level.m = entry.first; + level.u = entry.second; + level.m_estimated = true; + comparison.levels.push_back(level); + } + return comparison; + } + + // The same run three ways: shards only, table only, and both at once. + bool Run(const std::string& shards, cpplink::EdgeTable* table, unsigned threads, + cpplink::PredictReport* report, std::string* error) { + cpplink::PredictOptions options; + options.out_dir = shards.empty() ? "" : (dir_ / shards).string(); + options.table = table; + options.threads = threads; + return cpplink::Predict(*store_, comparisons_, plan_, scorer_, options, report, + error); + } + + static std::vector ReadShards(const std::string& directory) { + std::vector edges; + for (const auto& entry : std::filesystem::directory_iterator(directory)) { + std::ifstream file(entry.path(), std::ios::binary); + std::string bytes((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + for (size_t at = 8; at + 20 <= bytes.size(); at += 20) { + Edge edge; + std::memcpy(&edge.a, bytes.data() + at, 4); + std::memcpy(&edge.b, bytes.data() + at + 4, 4); + std::memcpy(&edge.gamma, bytes.data() + at + 8, 4); + std::memcpy(&edge.weight, bytes.data() + at + 12, 8); + edges.push_back(edge); + } + } + std::sort(edges.begin(), edges.end()); + return edges; + } + + // The partition with each cluster named by its lowest row, since which row + // union-find picks as representative depends on the order the edges came. + static std::vector Canonical(const cpplink::ClusterAssignment& assignment) { + std::vector lowest(assignment.root.size(), UINT32_MAX); + for (size_t row = 0; row < assignment.root.size(); ++row) { + lowest[assignment.root[row]] = + std::min(lowest[assignment.root[row]], static_cast(row)); + } + std::vector canonical(assignment.root.size()); + for (size_t row = 0; row < assignment.root.size(); ++row) { + canonical[row] = lowest[assignment.root[row]]; + } + return canonical; + } + + static std::vector Rows(const cpplink::EdgeTable& table) { + std::vector edges; + for (uint64_t i = 0; i < table.Size(); ++i) { + edges.push_back({table.a[i], table.b[i], table.gamma[i], table.weight[i]}); + } + std::sort(edges.begin(), edges.end()); + return edges; + } + + std::filesystem::path dir_; + cpplink::Schema schema_; + std::unique_ptr store_; + cpplink::BlockingPlan plan_; + cpplink::ComparisonSet comparisons_; + cpplink::Model model_; + cpplink::Scorer scorer_; +}; + +} // namespace + +TEST_F(EdgeTableFixture, TheTableHoldsWhatTheShardsHold) { + std::string error; + cpplink::PredictReport files; + ASSERT_TRUE(Run("shards", nullptr, 3, &files, &error)) << error; + const std::vector from_files = ReadShards((dir_ / "shards").string()); + ASSERT_FALSE(from_files.empty()); + + cpplink::EdgeTable table; + cpplink::PredictReport memory; + ASSERT_TRUE(Run("", &table, 3, &memory, &error)) << error; + EXPECT_TRUE(memory.shards.empty()); + EXPECT_EQ(memory.edges, files.edges); + EXPECT_EQ(table.Size(), files.edges); + EXPECT_EQ(Rows(table), from_files); + EXPECT_FALSE(std::filesystem::exists(dir_ / "shards2")); + + // Both at once: the same rows in both places. + cpplink::EdgeTable both; + ASSERT_TRUE(Run("shards2", &both, 2, &memory, &error)) << error; + EXPECT_EQ(Rows(both), ReadShards((dir_ / "shards2").string())); + EXPECT_EQ(Rows(both), from_files); + + // Nowhere to put them is refused. + cpplink::PredictReport nowhere; + EXPECT_FALSE(Run("", nullptr, 1, &nowhere, &error)); + EXPECT_NE(error.find("nowhere to put"), std::string::npos) << error; +} + +TEST_F(EdgeTableFixture, ClusteringTheTableIsClusteringTheShards) { + std::string error; + cpplink::PredictReport report; + cpplink::EdgeTable table; + ASSERT_TRUE(Run("shards", &table, 2, &report, &error)) << error; + + cpplink::ClusterOptions from_files; + from_files.edge_path = (dir_ / "shards").string(); + from_files.threshold = 0.0; + cpplink::ClusterAssignment files; + cpplink::ClusterReport files_report; + ASSERT_TRUE(cpplink::Cluster(*store_, from_files, &files, &files_report, &error)) + << error; + + cpplink::ClusterOptions from_table; + from_table.edges = &table; + from_table.threshold = 0.0; + cpplink::ClusterAssignment memory; + cpplink::ClusterReport memory_report; + ASSERT_TRUE(cpplink::Cluster(*store_, from_table, &memory, &memory_report, &error)) + << error; + EXPECT_EQ(Canonical(memory), Canonical(files)); + EXPECT_EQ(memory_report.edges_read, files_report.edges_read); + EXPECT_EQ(memory_report.edges_used, files_report.edges_used); + EXPECT_GT(memory.clusters, 0u); + + // A prediction table arriving as a stream: ids as text, weights as doubles, + // one column the reader does not need, and the rows in a different order. + cpplink::BatchBuilder builder; + const int id_a = builder.AddColumn("id_a", cpplink::ExportType::kString); + const int weight = builder.AddColumn("match_weight", cpplink::ExportType::kDouble); + const int id_b = builder.AddColumn("id_b", cpplink::ExportType::kString); + const int extra = builder.AddColumn("note", cpplink::ExportType::kString); + for (uint64_t i = table.Size(); i-- > 0;) { + builder.AppendString(id_a, store_->ids().Get(table.a[i])); + builder.AppendDouble(weight, table.weight[i]); + builder.AppendString(id_b, store_->ids().Get(table.b[i])); + builder.AppendString(extra, "x"); + } + ArrowSchema schema; + ArrowArray batch; + builder.ExportSchema(&schema); + ASSERT_TRUE(builder.ExportBatch(&batch, &error)) << error; + auto imported = arrow::ImportRecordBatch(&batch, &schema); + ASSERT_TRUE(imported.ok()) << imported.status().message(); + auto reader = arrow::RecordBatchReader::Make({*imported}); + ASSERT_TRUE(reader.ok()); + ArrowArrayStream stream; + ASSERT_TRUE(arrow::ExportRecordBatchReader(*reader, &stream).ok()); + + cpplink::ClusterOptions from_stream; + from_stream.stream = &stream; + from_stream.threshold = 0.0; + cpplink::ClusterAssignment streamed; + cpplink::ClusterReport streamed_report; + ASSERT_TRUE( + cpplink::Cluster(*store_, from_stream, &streamed, &streamed_report, &error)) + << error; + EXPECT_EQ(stream.release, nullptr) << "the stream is released"; + EXPECT_EQ(Canonical(streamed), Canonical(files)); + EXPECT_EQ(streamed_report.edges_used, files_report.edges_used); +} + +// The record ids leave as a dictionary over the id arena: the arena's own +// offsets and text as the dictionary, the rows as the indices, nothing copied. +TEST_F(EdgeTableFixture, IdsExportAsADictionaryOverTheArena) { + std::string error; + cpplink::PredictReport report; + auto table = std::make_shared(); + ASSERT_TRUE(Run("", table.get(), 2, &report, &error)) << error; + ASSERT_GT(table->Size(), 0u); + + const cpplink::IdColumn& ids = store_->ids(); + cpplink::BorrowedColumn id_a; + id_a.type = cpplink::ExportType::kDictionary; + id_a.length = static_cast(table->Size()); + id_a.values = table->a.data(); + id_a.dictionary_offsets = reinterpret_cast(ids.offsets.data()); + id_a.dictionary_text = ids.text.data(); + id_a.dictionary_size = static_cast(ids.offsets.size() - 1); + cpplink::BorrowedColumn weight; + weight.type = cpplink::ExportType::kDouble; + weight.length = id_a.length; + weight.values = table->weight.data(); + + cpplink::BatchBuilder builder; + builder.AddBorrowed("id_a", id_a); + builder.AddBorrowed("match_weight", weight); + const int note = builder.AddColumn("note", cpplink::ExportType::kString); + for (uint64_t i = 0; i < table->Size(); ++i) builder.AppendString(note, "n"); + builder.KeepAlive(table); + EXPECT_EQ(builder.Rows(), id_a.length); + + ArrowSchema schema; + ArrowArray batch; + builder.ExportSchema(&schema); + ASSERT_TRUE(builder.ExportBatch(&batch, &error)) << error; + auto imported = arrow::ImportRecordBatch(&batch, &schema); + ASSERT_TRUE(imported.ok()) << imported.status().message(); + const auto& rows = *imported; + ASSERT_TRUE(rows->ValidateFull().ok()) << rows->ValidateFull().message(); + EXPECT_EQ(rows->schema()->field(0)->type()->ToString(), + "dictionary"); + const auto& dictionary = static_cast(*rows->column(0)); + EXPECT_EQ(dictionary.dictionary()->length(), static_cast(kRecords)); + EXPECT_EQ(dictionary.dictionary()->data()->buffers[2]->data(), + reinterpret_cast(ids.text.data())) + << "the dictionary is the arena itself"; + const auto& values = + static_cast(*dictionary.dictionary()); + const auto& indices = static_cast(*dictionary.indices()); + for (int64_t i = 0; i < rows->num_rows(); ++i) { + EXPECT_EQ(values.GetString(indices.Value(i)), ids.Get(table->a[i])); + } + EXPECT_EQ(static_cast(*rows->column(1)).Value(0), + table->weight[0]); + + // The owned column moved out, so the builder is ragged until it is refilled; + // the borrowed ones are still there, and every live export holds the table. + ArrowArray again; + EXPECT_FALSE(builder.ExportBatch(&again, &error)); + EXPECT_NE(error.find("\"note\" has 0 rows"), std::string::npos) << error; + EXPECT_EQ(builder.Rows(), id_a.length); + EXPECT_GE(table.use_count(), 3) << "the test, the builder and the live export"; + cpplink::BatchBuilder borrowed_only; + borrowed_only.AddBorrowed("id_a", id_a); + borrowed_only.KeepAlive(table); + ASSERT_TRUE(borrowed_only.ExportBatch(&again, &error)) << error; + again.release(&again); + ASSERT_TRUE(borrowed_only.ExportBatch(&again, &error)) << "exported twice"; + EXPECT_EQ(again.children[0]->length, id_a.length); + again.release(&again); +} From d5f598212d314e88a07aa3469595dd0701f4080d Mon Sep 17 00:00:00 2001 From: 4ment Date: Mon, 21 Sep 2026 17:01:19 +1000 Subject: [PATCH 2/3] Include for std::isnan in the frame tests --- tests/arrow_export_test.cpp | 1 + tests/batch_loader_test.cpp | 1 + 2 files changed, 2 insertions(+) diff --git a/tests/arrow_export_test.cpp b/tests/arrow_export_test.cpp index 8d8d285..e53017b 100644 --- a/tests/arrow_export_test.cpp +++ b/tests/arrow_export_test.cpp @@ -8,6 +8,7 @@ #include "cpplink/arrow_export.hpp" +#include #include #include #include diff --git a/tests/batch_loader_test.cpp b/tests/batch_loader_test.cpp index 931654a..309f37d 100644 --- a/tests/batch_loader_test.cpp +++ b/tests/batch_loader_test.cpp @@ -11,6 +11,7 @@ #include "cpplink/batch_loader.hpp" +#include #include #include #include From c6b98977ac1b371144195a28b85fe06f2792d2df Mon Sep 17 00:00:00 2001 From: 4ment Date: Mon, 21 Sep 2026 17:53:09 +1000 Subject: [PATCH 3/3] Enhance documentation and examples for parquet support and predictions --- CONTRIBUTING.md | 4 ++ README.md | 10 ++-- docs/python.md | 74 +++++++++++++++++++++++++--- python/cpplink/__init__.py | 10 ++-- python/examples/historical_50k.ipynb | 23 ++++----- 5 files changed, 95 insertions(+), 26 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 7e7cfd7..35132e4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -51,6 +51,10 @@ python -m pytest python/tests The editable install redirects the Python sources to `python/cpplink/`, so a change there needs no rebuild; a change to `python/bindings/` or to the core needs the `pip install` line again. The binding is a second front end over the same stage functions, not a second implementation, and `python/tests/test_parity.py` holds it to the command line's output byte for byte. +The same package builds with `CPPLINK_WITH_ARROW=OFF`, which is what the wheels workflow ([.github/workflows/wheels.yml](.github/workflows/wheels.yml)) does through cibuildwheel and what CI's `no-arrow` job checks: the module then links nothing but the C++ standard library and a parquet file is read and written through pyarrow and pandas. +`python/tests/conftest.py` marks the tests that need the core to read parquet itself `needs_core_parquet`, and they are skipped in that build; a new test over a parquet path the core opens takes the same mark. +`test_reports.py` records a field count per bound report class, so a struct gaining a field the binding forgot is caught when the count is bumped deliberately. + ## Code style C++ follows the [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html) with two customizations, a 4-space indent and a 90-column limit, both in [.clang-format](.clang-format) and [CPPLINT.cfg](CPPLINT.cfg). diff --git a/README.md b/README.md index e612b9e..6c6abdb 100644 --- a/README.md +++ b/README.md @@ -101,11 +101,11 @@ The same pipeline from Python: ```python import cpplink -linker = cpplink.Linker("schema.json", "sample.parquet") +linker = cpplink.Linker("schema.json", "sample.parquet") # or a pandas frame model, report = linker.estimate(out="model.json") -linker.predict(model, "predictions.parquet", threshold=20) -result = linker.cluster("predictions.parquet", truth="sample.truth.csv") -print(result.quality) +predictions = linker.predict(model, threshold=20) # a pandas frame +clusters = linker.cluster(truth="sample.truth.csv") # another +print(linker.last_cluster.quality) ``` Between `init` and `estimate` sit the diagnostics that cost seconds and decide the quality of the result: `profile` says what each column can be worth before any model exists, `levels` places the fuzzy thresholds from the data, `explain-blocking` prices every source without enumerating a pair, and `recall` measures what blocking reaches. @@ -130,6 +130,8 @@ The Python package builds into the same environment, against the Arrow it holds: pip install -e . --no-build-isolation -Ccmake.define.CMAKE_PREFIX_PATH=$CONDA_PREFIX ``` +Outside conda it builds with no Arrow at all, `pip install . -Ccmake.define.CPPLINK_WITH_ARROW=OFF`, and reads and writes parquet through pandas instead; [From Python](https://4ment.github.io/cpplink/python/) says what that build leaves out. + Building with the tests, the sanitizer presets, the lint tools and the documentation are in [CONTRIBUTING.md](CONTRIBUTING.md). ## License diff --git a/docs/python.md b/docs/python.md index df1e1a9..1fc1e47 100644 --- a/docs/python.md +++ b/docs/python.md @@ -13,7 +13,7 @@ A worked example over a real dataset, with plots and tables at every stage, is t ## Install -Inside the project's conda environment: +Inside the project's conda environment, against the Arrow C++ it holds: ```sh conda env update -f environment.yml @@ -25,11 +25,27 @@ python -c "import cpplink; print(cpplink.__version__)" The build goes through [scikit-build-core](https://scikit-build-core.readthedocs.io/) and lands in `build/python/`, beside the command line's `build/`. The editable install redirects the Python sources to `python/cpplink/`, so a change there needs no rebuild; a change to the bindings or the core does, with the same `pip install` line. +Outside conda, in any environment with a C++17 compiler, the package builds with no Arrow at all: + +```sh +pip install . -Ccmake.define.CPPLINK_WITH_ARROW=OFF +``` + +That is the build the wheels workflow makes, and it links nothing but the C++ standard library: `pandas`, `numpy` and `pyarrow` are its only dependencies, at whatever versions the environment holds. +`cpplink.parquet_supported()` says which build is loaded. + !!! note "What the module links" - Nothing but the C++ standard library. Every frame, table and file crosses between Python and the core through the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html), three C structs with a frozen ABI that the core vendors, so no Arrow C++ object and no `pyarrow` version is part of the contract. `pyarrow` is what pandas uses to export a frame and what reads a parquet file into the module, and any version of it will do. - The conda build above links Arrow C++ as well, so that `predict(out="x.parquet")` and the command line write parquet themselves; a wheel built with `-DCPPLINK_WITH_ARROW=OFF` leaves that to pandas. + The conda build links Arrow C++ as well, so that `predict(out="x.parquet")` and the command line write parquet themselves. + +What a build without Arrow does differently, all of it at the edges: + +- A parquet path given to `Linker`, `init`, `cluster` or `gen_sample` is read or written through `pyarrow` and `pandas`, from the same rows the conda build reads and writes itself, so the model, the predictions and the clusters are the same. +- `cpplink.run` and `python -m cpplink` still dispatch every command, but the ones that open a parquet file fail with the reason, since that is the core reading the file. +- `cluster_file`, which loads the id column alone, and `merge_predictions` over binary shards, which needs the input to turn rows back into ids, need the core to read parquet and are refused the same way. + +The test suite marks those paths `needs_core_parquet` and skips them there; everything else runs on both builds. ## The command line, in process @@ -55,6 +71,14 @@ frame = cpplink.gen_sample(rows=200000, truth="sample.truth.csv") # a pandas f cpplink.gen_sample("sample.parquet", rows=200000) # or the file ``` +`seed` and `duplicate_rate` are the command's flags. +`out_b` makes the link fixture, originals in the first file and every planted duplicate in a later one: a list of paths with `out`, or a count of later files without it, which then returns a list of frames. + +```python +cpplink.gen_sample("a.parquet", rows=20000, truth="link.truth.csv", out_b=["b.parquet"]) +a, b = cpplink.gen_sample(rows=20000, out_b=1) +``` + ### 2. Draft a schema and load ```python @@ -65,6 +89,10 @@ linker = cpplink.Linker(schema, frame) print(linker.inspect()) # cardinality and memory per column ``` +`init` reads no rows: a parquet path gives up its footer, and a frame, a pyarrow table or schema, or a `{name: input}` dict of them gives up its Arrow schema, which says what the footer would have. +`roles={"birth_place": "city"}` overrides a guessed role, as `--role COLUMN=ROLE` does, and `id=` names the record id; the roles `init` accepts are `cpplink.known_roles()`. +The report has a `DraftColumn` per column, with its `type`, its `role` and whether the role was `role_given` or guessed. + `Linker` takes a `Schema` or the path of one, and the data: a pandas frame, a pyarrow table or reader, a polars frame or anything else with `__arrow_c_stream__`, or a parquet path. Several inputs are a `{name: input}` dict, each becoming a dataset named by its key, or a list of paths named by their stems; two inputs and no `mode` means linking them. @@ -80,14 +108,26 @@ The frame's index is not read; the record id must be a column, and an integer id Several files are several datasets and default to linking; `mode="dedup"` or `"link-and-dedup"` scores the within-file pairs too, exactly as `--mode` does. `all_pairs=True` is `--all-pairs`. -A `Schema` is edited by dict round trip, which is what `init` itself does: +What was loaded is on the `Linker` itself: `records`, `datasets`, `mode` (the pairs the store admits: `all pairs` deduplicating, `cross-dataset pairs` linking), `all_pairs`, `files` and `stats`, the `LoadStats` with the rows and seconds per input. +`id_of(row)` and `row_of(id)` go between a row index and a record id, and `repr(linker)` says what it holds. + +```python +>>> linker +Linker(200,000 records, 1 dataset, mode=all pairs) +``` + +A `Schema` comes from `Schema.from_file`, `Schema.from_json` or `Schema.from_dict`, and is edited by dict round trip, which is what `init` itself does: ```python d = schema.to_dict() d["blocking"].append({"type": "exact_value", "column": "postcode"}) schema = cpplink.Schema.from_dict(d) +schema.save("schema.json") ``` +`to_json` is the writer `cpplink init` uses, and `save` writes the source text where the schema came from one. +The read-only side describes what was parsed: `unique_id`, `column_names`, `input_columns` (the id and every column that is not derived, which is what a reader is asked for), `comparison_names`, `comparisons` (the description of each level, per comparison), `blocking` and `gamma_width`. + ### 3. Price the blocking ```python @@ -114,8 +154,12 @@ model.lambda_, model.prior_weight() [(c.name, [(l.label, l.m, l.u) for l in c.levels]) for c in model.comparisons] ``` -`estimate` takes every flag of the command by name: `u_sample`, `session_pairs`, `threads`, `iterations`, `lambda_`, `seed`, `fuzzy_u`, `ball_budget`, `tie_holdout`, `tied_bits`, `interactions`, `max_interactions`, `interaction_bits`, and so on. -An `EstimateOptions` can be built and passed as `options=` instead. +`estimate` takes every flag of the command by name: `u_sample`, `session_pairs`, `threads`, `iterations`, `lambda_`, `seed`, `fuzzy_u`, `ball_budget`, `tie_holdout`, `tied_bits`, `tie_sample_rows`, `interactions`, `max_interactions`, `interaction_bits`, `interaction_clamp`, `min_interaction_sessions` and `min_pairs_per_parameter`. +An `EstimateOptions` can be built and passed as `options=` instead, with the keywords applied over it. + +A `Model` is the `model.json` file as an object: `Model.from_file` and `Model.from_json` read one, `to_json`, `to_dict` and `save` write it, and `text` or `repr(model)` is the table `estimate` prints. +It carries `lambda_`, `lambda_basis`, `records`, a `ModelComparison` per comparison with a `ModelLevel` per level (`label`, `m`, `u` and `weight()`, the bits `log2(m / u)`), and a `ModelInteraction` per fitted two-way term. +Every method that takes a `model` takes a `Model` or the path of one. ### 6. Score the candidates @@ -139,6 +183,16 @@ table = pa.table(report.table) # zero copy; the ids are a dictionary over th ``` The frame decodes that dictionary into plain string columns, since a pandas `Categorical` over every record of the store is neither cheap nor what a frame of predictions wants; the decode copies only the ids the rows name. +`cpplink.to_pandas(table)` is that decode, for any table the core hands out. + +`cpplink.weight_for_probability` and `cpplink.probability_for_weight` are the two conversions the reports use, so a threshold can be stated either way and read back: + +```python +>>> cpplink.probability_for_weight(20) +0.9999990463265931 +>>> cpplink.weight_for_probability(0.99) +6.629356620079609 +``` ### 7. Join the predictions into clusters @@ -167,12 +221,18 @@ linker.profile() # what the columns are worth, befor linker.levels() # (LevelsReport, Schema with the proposal) linker.simplify(model, min_gap=1.0) # (SimplifyReport, Schema or None) linker.completeness(model) # blocking recall with no truth file -linker.rescore(model, "spill/", "again.parquet", threshold=25) +linker.rescore(model, "spill/", threshold=25) # a frame, with the report on last_rescore cpplink.merge_predictions("shards/", "merged.csv", schema="schema.json", files=["sample.parquet"]) +cpplink.cluster_file("schema.json", "sample.parquet", "predictions.parquet", truth="sample.truth.csv") ``` `explain` resolves ids the way the command does, including `dataset:id` where the inputs share ids; `by_row=True` takes row indices. +The `Explanation` has `levels`, comparison name to the label the pair landed on, and with a model a `waterfall`, the `PairWaterfall` ledger of `prior`, one `WaterfallStep` per comparison with its `bits`, `tf` move and `running` total, the `interactions`, and the final `weight` and `probability`; `json()` is the command's `--json`. Its `weight` is the scorer's own, and the parity test asserts it equals the weight in the predictions file. +`explain` takes `fuzzy_tf` and `interactions=False` for the same reason the command does: scored under different options than the run, it explains a different weight than the file holds. + +`rescore` replays a spill `predict` wrote under a new model and returns the predictions as `predict` does, with `out=` writing them the same two ways. +`levels` and `simplify` return the schema with the proposal written in, or `None` from `simplify` where nothing merged, and `out=` writes it. ## Reports diff --git a/python/cpplink/__init__.py b/python/cpplink/__init__.py index 53b8475..613c6c0 100644 --- a/python/cpplink/__init__.py +++ b/python/cpplink/__init__.py @@ -3,10 +3,10 @@ """cpplink from Python: the command line in process, and a ``Linker`` over it. Two ways in. ``run(args)`` forwards an argument list to the command line -dispatcher exactly as the binary would receive it. ``Linker`` loads the parquet -input once and runs every stage against the same store, returning the reports -as objects rather than text. Both call the same core; no number is computed in -Python. +dispatcher exactly as the binary would receive it. ``Linker`` loads the input +once, a data frame or a parquet file, and runs every stage against the same +store, returning predictions and clusters as frames and the reports as objects +rather than text. Both call the same core; no number is computed in Python. """ from __future__ import annotations @@ -58,6 +58,7 @@ SimplifyReport, __version__, known_roles, + parquet_supported, probability_for_weight, run, weight_for_probability, @@ -108,6 +109,7 @@ "init", "known_roles", "merge_predictions", + "parquet_supported", "probability_for_weight", "run", "to_dict", diff --git a/python/examples/historical_50k.ipynb b/python/examples/historical_50k.ipynb index b663db2..5e72e1a 100644 --- a/python/examples/historical_50k.ipynb +++ b/python/examples/historical_50k.ipynb @@ -2289,7 +2289,7 @@ "source": [ "## 8. Predict: score every candidate pair\n", "\n", - "`predict` walks the blocked union once, one thread per core, and writes every pair whose weight clears the threshold.\n", + "`predict` walks the blocked union once, one thread per core, and returns every pair whose weight clears the threshold as a frame, written to the file named as well so the sweep below can re-read it.\n", "Three prunes keep the metrics off most pairs and are all exact: a per-value signature bounds a string metric before a character is read, the *ceiling* grants every comparison the best level the cheap bounds still admit and drops the pair when even that sum is under the threshold, and a per-pattern bracket on the term-frequency adjustment settles most patterns without a table lookup.\n", "The threshold is 0 bits here, the posterior at even odds, so the whole frontier can be read off one file." ] @@ -2339,9 +2339,10 @@ ], "source": [ "t = time.time()\n", - "predict = linker.predict(model, WORK / \"predictions.parquet\", threshold=0)\n", + "preds = linker.predict(model, WORK / \"predictions.parquet\", threshold=0)\n", "timings[\"predict\"] = time.time() - t\n", - "predict_x = linker.predict(model_x, WORK / \"predictions_interactions.parquet\", threshold=0)\n", + "predict = linker.last_predict\n", + "preds_x = linker.predict(model_x, WORK / \"predictions_interactions.parquet\", threshold=0)\n", "print(predict)" ] }, @@ -2375,14 +2376,13 @@ } ], "source": [ - "def read_predictions(path):\n", - " df = pq.read_table(path).to_pandas()\n", + "def mark_truth(df):\n", " key = np.where(df.id_a < df.id_b, df.id_a + \"|\" + df.id_b, df.id_b + \"|\" + df.id_a)\n", " df[\"match\"] = pd.Series(key).isin({f\"{a}|{b}\" for a, b in truth_pairs}).values\n", " return df\n", "\n", - "preds = read_predictions(WORK / \"predictions.parquet\")\n", - "preds_x = read_predictions(WORK / \"predictions_interactions.parquet\")\n", + "preds = mark_truth(preds)\n", + "preds_x = mark_truth(preds_x)\n", "\n", "fig, ax = plt.subplots(figsize=(10, 4))\n", "bins = np.arange(0, 82, 2)\n", @@ -2557,7 +2557,8 @@ "def sweep(path, thresholds):\n", " rows = []\n", " for th in thresholds:\n", - " q = linker.cluster(path, threshold=th, truth=truth_path).quality\n", + " linker.cluster(path, threshold=th, truth=truth_path)\n", + " q = linker.last_cluster.quality\n", " rows.append({\"threshold\": th, \"precision\": q.precision, \"recall\": q.recall, \"f1\": q.f1})\n", " return pd.DataFrame(rows)\n", "\n", @@ -2654,8 +2655,9 @@ } ], "source": [ - "result = linker.cluster(WORK / \"predictions_interactions.parquet\", threshold=best_x.threshold,\n", - " truth=truth_path, out=WORK / \"clusters.csv\")\n", + "clusters = linker.cluster(WORK / \"predictions_interactions.parquet\", threshold=best_x.threshold,\n", + " truth=truth_path, out=WORK / \"clusters.csv\")\n", + "result = linker.last_cluster\n", "print(result)" ] }, @@ -2689,7 +2691,6 @@ } ], "source": [ - "clusters = pd.read_csv(WORK / \"clusters.csv\")\n", "found_sizes = clusters.groupby(\"cluster_id\").size()\n", "edges = [1, 2, 3, 4, 5, 6, 11, 21, 51, 1000]\n", "labels = [\"1\", \"2\", \"3\", \"4\", \"5\", \"6-10\", \"11-20\", \"21-50\", \"51+\"]\n",