pqbench 0.0.2__tar.gz

pqbench-0.0.2/Cargo.toml

@@ -0,0 +1,42 @@
+[workspace]
+resolver = "2"
+members = [
+    "crates/aipnaming",
+    "crates/aipnaming-cli",
+    "crates/docscheck",
+    "crates/docscheck-cli",
+    "crates/pqbench",
+    "crates/pqbench-cli",
+]
+exclude = ["python"]
+default-members = [
+    "crates/aipnaming",
+    "crates/aipnaming-cli",
+    "crates/docscheck",
+    "crates/docscheck-cli",
+    "crates/pqbench",
+    "crates/pqbench-cli",
+]
+
+[workspace.package]
+version = "0.0.2"
+edition = "2021"
+rust-version = "1.91.1"
+license = "MIT OR Apache-2.0"
+authors = ["pqbench contributors"]
+repository = "https://github.com/pqbench/pqbench"
+
+[workspace.dependencies]
+aipnaming = { version = "0.0.2", path = "crates/aipnaming" }
+docscheck = { version = "0.0.2", path = "crates/docscheck" }
+pqbench = { version = "0.0.2", path = "crates/pqbench" }
+
+[profile.dev]
+opt-level = 0
+debug = 1
+incremental = true
+
+# The profile that 'dist' will build with
+[profile.dist]
+inherits = "release"
+lto = "thin"

pqbench-0.0.2/PKG-INFO

@@ -0,0 +1,188 @@
+Metadata-Version: 2.4
+Name: pqbench
+Version: 0.0.2
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: License :: OSI Approved :: MIT License
+Classifier: License :: OSI Approved :: Apache Software License
+Summary: Python bindings for the pqbench command line.
+Keywords: parquet,compression,benchmark
+Author: pqbench contributors
+License-Expression: MIT OR Apache-2.0
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Documentation, https://github.com/pqbench/pqbench/tree/main/python
+Project-URL: Homepage, https://github.com/pqbench/pqbench
+Project-URL: Repository, https://github.com/pqbench/pqbench
+
+# pqbench
+
+PyO3 bindings for [pqbench](https://github.com/pqbench/pqbench). `pip install`
+builds a native wheel; a separate `pqbench` binary is not required. Two
+surfaces, one wheel:
+
+* **`pqbench.native`** — the commands: one function per library path, with the
+  same arguments. Every name is the Rust function of the same path.
+* **`pqbench`** — the ergonomic surface over them: the client handles every
+  command takes (`Config`, `Dialect`, `Lake`, `Layer`, `Workspace`) and `walk`,
+  the whole-metastore walk composed once.
+
+## Install
+
+```sh
+pip install pqbench
+```
+
+The wheel carries every storage backend (`delta-s3`, `iceberg-s3`), so `s3://`
+reads and Delta/Iceberg tables work out of the box. A storage read
+authenticates with the `credentials=` argument (a `Credentials` record) or,
+for `bytemass`, with `env` on its request — or with the process environment.
+
+From a checkout:
+
+```sh
+pip install maturin
+maturin develop --manifest-path python/Cargo.toml
+```
+
+## The commands
+
+Each command is the library function of the same path, under `pqbench.native`,
+and takes the handles and the resource's name like the Rust signature: the name
+is the AIP resource name a parent record carries, so a record feeds the next
+command as it is. `credentials.get` takes a table's name and returns its lease,
+the `Credentials` record a storage read takes, and a record is its model's
+class: `record["uri"]`, `record.uri` and `record.asdict()` agree.
+
+They are the walk's own pipeline — metastore, catalog, schema, table,
+credentials, partition, file, bytemass — and they stop at `bytemass`: the
+commands downstream of a byte-mass stream (`viz`, `diff`, `profile`,
+`experiment`), the standalone codec benchmarks (`lz`, `compression`) and
+`skill` are separate tools, and are not bound.
+
+```python
+import asyncio
+import pqbench
+import pqbench.native
+
+async def main():
+    workspace = pqbench.Workspace(pqbench.Config("https://…", token="…"))
+    lake = pqbench.Lake(workspace)
+
+    # The walk, one function per resource. The I/O verbs are awaitable.
+    await pqbench.native.metastore.get.get(workspace)
+    await pqbench.native.catalog.get.get(workspace, "catalogs/dbx_samples")
+    await pqbench.native.catalog.list.list(workspace)
+    await pqbench.native.schema.list.list(workspace, "catalogs/dbx_samples")
+    refs = await pqbench.native.table.list.list(
+        workspace, "catalogs/dbx_samples/schemas/nyctaxi"
+    )
+    await pqbench.native.credentials.check.check(
+        workspace, "catalogs/dbx_samples/schemas/nyctaxi/tables/trips"
+    )
+    # `credentials.get.get` is name -> lease: the lease is what a storage read
+    # takes, from the table's name alone.
+    credentials = await pqbench.native.credentials.get.get(workspace, refs[0]["name"])
+    table = await pqbench.native.table.get.get(
+        workspace,
+        lake,
+        refs[0]["name"], refs[0].get("uri"), credentials,
+    )
+
+    # `partition.list.list` and `file.list.list` stream (async iterators); the
+    # single-value verbs are awaitable.
+    async for partition in pqbench.native.partition.list.list(
+        lake, table["name"], table["uri"], table.get("credentials"), 86_400_000
+    ):
+        async for file in pqbench.native.file.list.list(
+            lake, partition["uri"], partition.get("credentials"), partition["commit_versions"]
+        ):
+            # The pipeline stops here: `bytemass` measures the file's bytes.
+            rows = await pqbench.native.bytemass.bytemass(
+                lake, pqbench.native.BytemassRequest(inputs=[file["uri"]])
+            )
+            print(len(rows), file["uri"])
+
+asyncio.run(main())
+```
+
+### The walk
+
+`walk` is the one composition rather than a binding: it spans the catalog
+levels and descends every table the principal may read outside compute, with
+the tables in flight together (bounded by `concurrency`, 32 by default), and
+yields each table file.
+
+```python
+async for file in pqbench.walk(workspace):
+    print(file["uri"])
+```
+
+`catalog=` descends one catalog instead of every readable one. A table that
+cannot be read is skipped (logged at `debug`), so a mixed metastore still
+walks; records stream, so the walk holds no more than a bounded queue of them.
+
+The Unity Catalog API limits requests per endpoint and per workspace, so its
+calls are paced at `requests_per_second` — 15, one bucket per endpoint, and `0`
+turns pacing off. The storage reads are not paced: object storage has no such
+limit and retries its own. Anything skipped is counted in a `warning` when the
+walk ends, so a rate limit shows up as a count rather than as silently missing
+files.
+
+`bytemass` is the last command: it reads what `file.list.list` yields — a
+path or URI per file — and returns one `MassRow` per column chunk.
+
+```python
+rows = await pqbench.native.bytemass.bytemass(
+    lake, pqbench.native.BytemassRequest(inputs=["s3://…/part-00000.parquet"])
+)
+```
+
+[`samples/walk.py`](samples/walk.py) is a runnable walk over the **whole
+metastore**: it calls `walk`, spans every catalog and schema, descends the one
+catalog the principal may read outside compute, and prints each table file.
+
+```sh
+PQB_ENDPOINT=https://<workspace>.cloud.databricks.com PQB_TOKEN=<bearer> \
+  AWS_REGION=us-east-2 python samples/walk.py
+```
+
+## Configuration
+
+`Config(endpoint, token=None, dialect=Dialect.UNITY)` carries the endpoint, the
+bearer, and the dialect (`UNITY` or `ICEBERG`); `Layer(endpoint=None,
+token=None, dialect=None).resolve()` combines configuration layers into a
+`Config`.
+
+`credentials` is a `Credentials` record — what `credentials.get` returns — or
+the mapping of the same fields; a read given none falls back to the process
+environment. A request's `env` is the library's `BTreeMap<String, String>`,
+passed straight through. A `BytemassRequest` names `inputs` (paths or glob
+masks), `env` (the `AWS_*` options a remote read uses) and `indexes` (load the
+parquet page indexes too, one extra range per file).
+
+## Publish to PyPI
+
+[`.github/workflows/publish-pypi.yml`](../.github/workflows/publish-pypi.yml)
+builds manylinux, macOS, and Windows wheels with
+[maturin](https://www.maturin.rs/) and uploads them with
+[PyPI trusted publishing](https://docs.pypi.org/trusted-publishers/) — no token
+is stored. It has its own dispatch, after the matching tag is released:
+`gh workflow run publish-pypi.yml`, or **Publish Python package** from the
+Actions tab. Publishing is separate because PyPI names the workflow file that
+may upload, and its trusted publisher does not accept a job defined in a
+workflow that another workflow calls.
+
+Its version is the workspace's: `scripts/version.sh sync` writes
+`[workspace.package] version` from `Cargo.toml` into `python/Cargo.toml` before
+maturin reads it, and `make check-version` fails when the two disagree —
+`python/` is its own crate and cannot inherit it.
+
+Before the first upload:
+
+1. On PyPI, add a pending publisher for this repository: owner `pqbench`,
+   project `pqbench`, workflow `publish-pypi.yml`, environment `pypi`.
+2. In the GitHub repository settings, create an environment named `pypi`.
+3. Run the **Publish Python package** workflow.
+

pqbench-0.0.2/README.md

@@ -0,0 +1,169 @@
+# pqbench
+
+PyO3 bindings for [pqbench](https://github.com/pqbench/pqbench). `pip install`
+builds a native wheel; a separate `pqbench` binary is not required. Two
+surfaces, one wheel:
+
+* **`pqbench.native`** — the commands: one function per library path, with the
+  same arguments. Every name is the Rust function of the same path.
+* **`pqbench`** — the ergonomic surface over them: the client handles every
+  command takes (`Config`, `Dialect`, `Lake`, `Layer`, `Workspace`) and `walk`,
+  the whole-metastore walk composed once.
+
+## Install
+
+```sh
+pip install pqbench
+```
+
+The wheel carries every storage backend (`delta-s3`, `iceberg-s3`), so `s3://`
+reads and Delta/Iceberg tables work out of the box. A storage read
+authenticates with the `credentials=` argument (a `Credentials` record) or,
+for `bytemass`, with `env` on its request — or with the process environment.
+
+From a checkout:
+
+```sh
+pip install maturin
+maturin develop --manifest-path python/Cargo.toml
+```
+
+## The commands
+
+Each command is the library function of the same path, under `pqbench.native`,
+and takes the handles and the resource's name like the Rust signature: the name
+is the AIP resource name a parent record carries, so a record feeds the next
+command as it is. `credentials.get` takes a table's name and returns its lease,
+the `Credentials` record a storage read takes, and a record is its model's
+class: `record["uri"]`, `record.uri` and `record.asdict()` agree.
+
+They are the walk's own pipeline — metastore, catalog, schema, table,
+credentials, partition, file, bytemass — and they stop at `bytemass`: the
+commands downstream of a byte-mass stream (`viz`, `diff`, `profile`,
+`experiment`), the standalone codec benchmarks (`lz`, `compression`) and
+`skill` are separate tools, and are not bound.
+
+```python
+import asyncio
+import pqbench
+import pqbench.native
+
+async def main():
+    workspace = pqbench.Workspace(pqbench.Config("https://…", token="…"))
+    lake = pqbench.Lake(workspace)
+
+    # The walk, one function per resource. The I/O verbs are awaitable.
+    await pqbench.native.metastore.get.get(workspace)
+    await pqbench.native.catalog.get.get(workspace, "catalogs/dbx_samples")
+    await pqbench.native.catalog.list.list(workspace)
+    await pqbench.native.schema.list.list(workspace, "catalogs/dbx_samples")
+    refs = await pqbench.native.table.list.list(
+        workspace, "catalogs/dbx_samples/schemas/nyctaxi"
+    )
+    await pqbench.native.credentials.check.check(
+        workspace, "catalogs/dbx_samples/schemas/nyctaxi/tables/trips"
+    )
+    # `credentials.get.get` is name -> lease: the lease is what a storage read
+    # takes, from the table's name alone.
+    credentials = await pqbench.native.credentials.get.get(workspace, refs[0]["name"])
+    table = await pqbench.native.table.get.get(
+        workspace,
+        lake,
+        refs[0]["name"], refs[0].get("uri"), credentials,
+    )
+
+    # `partition.list.list` and `file.list.list` stream (async iterators); the
+    # single-value verbs are awaitable.
+    async for partition in pqbench.native.partition.list.list(
+        lake, table["name"], table["uri"], table.get("credentials"), 86_400_000
+    ):
+        async for file in pqbench.native.file.list.list(
+            lake, partition["uri"], partition.get("credentials"), partition["commit_versions"]
+        ):
+            # The pipeline stops here: `bytemass` measures the file's bytes.
+            rows = await pqbench.native.bytemass.bytemass(
+                lake, pqbench.native.BytemassRequest(inputs=[file["uri"]])
+            )
+            print(len(rows), file["uri"])
+
+asyncio.run(main())
+```
+
+### The walk
+
+`walk` is the one composition rather than a binding: it spans the catalog
+levels and descends every table the principal may read outside compute, with
+the tables in flight together (bounded by `concurrency`, 32 by default), and
+yields each table file.
+
+```python
+async for file in pqbench.walk(workspace):
+    print(file["uri"])
+```
+
+`catalog=` descends one catalog instead of every readable one. A table that
+cannot be read is skipped (logged at `debug`), so a mixed metastore still
+walks; records stream, so the walk holds no more than a bounded queue of them.
+
+The Unity Catalog API limits requests per endpoint and per workspace, so its
+calls are paced at `requests_per_second` — 15, one bucket per endpoint, and `0`
+turns pacing off. The storage reads are not paced: object storage has no such
+limit and retries its own. Anything skipped is counted in a `warning` when the
+walk ends, so a rate limit shows up as a count rather than as silently missing
+files.
+
+`bytemass` is the last command: it reads what `file.list.list` yields — a
+path or URI per file — and returns one `MassRow` per column chunk.
+
+```python
+rows = await pqbench.native.bytemass.bytemass(
+    lake, pqbench.native.BytemassRequest(inputs=["s3://…/part-00000.parquet"])
+)
+```
+
+[`samples/walk.py`](samples/walk.py) is a runnable walk over the **whole
+metastore**: it calls `walk`, spans every catalog and schema, descends the one
+catalog the principal may read outside compute, and prints each table file.
+
+```sh
+PQB_ENDPOINT=https://<workspace>.cloud.databricks.com PQB_TOKEN=<bearer> \
+  AWS_REGION=us-east-2 python samples/walk.py
+```
+
+## Configuration
+
+`Config(endpoint, token=None, dialect=Dialect.UNITY)` carries the endpoint, the
+bearer, and the dialect (`UNITY` or `ICEBERG`); `Layer(endpoint=None,
+token=None, dialect=None).resolve()` combines configuration layers into a
+`Config`.
+
+`credentials` is a `Credentials` record — what `credentials.get` returns — or
+the mapping of the same fields; a read given none falls back to the process
+environment. A request's `env` is the library's `BTreeMap<String, String>`,
+passed straight through. A `BytemassRequest` names `inputs` (paths or glob
+masks), `env` (the `AWS_*` options a remote read uses) and `indexes` (load the
+parquet page indexes too, one extra range per file).
+
+## Publish to PyPI
+
+[`.github/workflows/publish-pypi.yml`](../.github/workflows/publish-pypi.yml)
+builds manylinux, macOS, and Windows wheels with
+[maturin](https://www.maturin.rs/) and uploads them with
+[PyPI trusted publishing](https://docs.pypi.org/trusted-publishers/) — no token
+is stored. It has its own dispatch, after the matching tag is released:
+`gh workflow run publish-pypi.yml`, or **Publish Python package** from the
+Actions tab. Publishing is separate because PyPI names the workflow file that
+may upload, and its trusted publisher does not accept a job defined in a
+workflow that another workflow calls.
+
+Its version is the workspace's: `scripts/version.sh sync` writes
+`[workspace.package] version` from `Cargo.toml` into `python/Cargo.toml` before
+maturin reads it, and `make check-version` fails when the two disagree —
+`python/` is its own crate and cannot inherit it.
+
+Before the first upload:
+
+1. On PyPI, add a pending publisher for this repository: owner `pqbench`,
+   project `pqbench`, workflow `publish-pypi.yml`, environment `pypi`.
+2. In the GitHub repository settings, create an environment named `pypi`.
+3. Run the **Publish Python package** workflow.

pqbench-0.0.2/crates/pqbench/Cargo.toml

@@ -0,0 +1,53 @@
+[package]
+name = "pqbench"
+version.workspace = true
+edition.workspace = true
+rust-version.workspace = true
+license.workspace = true
+authors.workspace = true
+repository.workspace = true
+description = "A benchmarking library for Parquet files (in-memory read/write throughput, compression ratio)."
+
+[features]
+default = []
+# S3 object reads. Without it, `s3://` fails at runtime, not at compile time.
+# reqwest signs a HEAD to read `x-amz-storage-class`, which object_store drops.
+aws = ["dep:object_store", "object_store/aws", "dep:async-trait"]
+# Delta snapshot resolution requires Rust 1.94.1 or newer.
+delta = ["dep:deltalake", "dep:futures"]
+delta-s3 = ["delta", "deltalake/s3", "aws"]
+# Iceberg snapshot planning reads metadata JSON and Avro manifests.
+iceberg = ["dep:apache-avro"]
+iceberg-s3 = ["iceberg", "aws"]
+
+[dependencies]
+snappy_src = "0.2.5"
+zstd = "0.13"
+lz4 = "1"
+flate2 = { version = "1", default-features = false, features = ["zlib"] }
+glob = "0.3"
+# Codecs are enabled so `read_sample` can decode compressed files; without
+# them parquet-rs fails at runtime (`Disabled feature at compile time`).
+parquet = { version = "59", default-features = false, features = ["snap", "brotli", "lz4", "zstd", "flate2-rust_backend"] }
+# Native roots only: the OS certificate store, where internal/private CAs live.
+# `rustls-tls-native-roots` also brings the rustls backend and the ring
+# provider, so the client still has TLS without `aws` (object_store is what
+# otherwise enables them).
+reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "json"] }
+async-trait = { version = "0.1", optional = true }
+bytes = "1"
+parquet-format-safe = "0.2.4"
+serde = { version = "1.0.229", features = ["derive"] }
+serde_json = "1.0.151"
+deltalake = { version = "1.0", default-features = false, features = ["rustls"], optional = true }
+futures = { version = "0.3", optional = true }
+futures-util = { version = "0.3", default-features = false, features = ["alloc"] }
+url = "2"
+# Must match delta-rs's ObjectStore trait version: stores cross that internal
+# boundary. Keep the pin until delta-rs upgrades its dependency.
+object_store = { version = "=0.13.2", optional = true }
+apache-avro = { version = "0.17", optional = true }
+
+[dev-dependencies]
+tempfile = "3"
+tokio = { version = "1", features = ["rt", "macros"] }

pqbench-0.0.2/crates/pqbench/src/bench.rs

@@ -0,0 +1,112 @@
+//! Shared codec-sweep mechanism for the `lz` and `compression` commands.
+//!
+//! The two commands own their own request types ([`crate::lz::LzRequest`],
+//! [`crate::compression::CompressionRequest`]); this module holds only the
+//! resolution they both need: parsing `codec@level` specs, the analytics
+//! config, and the raw pass count. It deliberately carries no request type of
+//! its own, so neither command depends on the other. Its [`Error`] is the error
+//! type both sweeps and their renderers return.
+
+use crate::codecs::{Codec, CodecImpl, Error as CodecError};
+use crate::stats;
+
+/// The measurement decisions a request resolves to: the codec×level set, the
+/// analytics config, and the raw pass count the upstream wants.
+pub(crate) struct Sweep {
+    pub(crate) codec_configs: Vec<(Codec, u8)>,
+    pub(crate) stats_config: stats::Config,
+    pub(crate) passes: u32,
+}
+
+/// Resolve the shared sweep decisions from a command's arguments.
+pub(crate) fn resolve(
+    codec_specs: &[String],
+    samples: u32,
+    warmup_iterations: u32,
+    mode: stats::Mode,
+) -> Result<Sweep, Error> {
+    Ok(Sweep {
+        codec_configs: parse_codec_configs(codec_specs)?,
+        stats_config: stats::Config {
+            warmup_iterations: warmup_iterations as usize,
+            mode,
+        },
+        passes: samples + warmup_iterations,
+    })
+}
+
+/// Errors from a codec sweep: spec parsing, the codec/parquet layers, or
+/// rendering.
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum Error {
+    /// A `codec@level` spec could not be parsed.
+    Spec(String),
+    /// The codec layer failed.
+    Codec(CodecError),
+    /// The parquet layer failed.
+    Parquet(crate::third_party::parquet::api::Error),
+    /// A renderer could not serialize its output.
+    Json(serde_json::Error),
+}
+
+impl std::fmt::Display for Error {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Error::Spec(e) => write!(f, "{e}"),
+            Error::Codec(e) => write!(f, "{e}"),
+            Error::Parquet(e) => write!(f, "{e}"),
+            Error::Json(e) => write!(f, "json: {e}"),
+        }
+    }
+}
+
+impl std::error::Error for Error {}
+
+impl From<CodecError> for Error {
+    fn from(e: CodecError) -> Self {
+        Error::Codec(e)
+    }
+}
+
+impl From<crate::third_party::parquet::api::Error> for Error {
+    fn from(e: crate::third_party::parquet::api::Error) -> Self {
+        Error::Parquet(e)
+    }
+}
+
+impl From<serde_json::Error> for Error {
+    fn from(e: serde_json::Error) -> Self {
+        Error::Json(e)
+    }
+}
+
+/// Parse the `codec@level` specs; empty selects every wired codec at its first
+/// level.
+fn parse_codec_configs(specs: &[String]) -> Result<Vec<(Codec, u8)>, Error> {
+    if specs.is_empty() {
+        return Ok(Codec::all().map(|c| (c, default_level(c))).collect());
+    }
+    specs.iter().map(|spec| parse_spec(spec)).collect()
+}
+
+/// One `codec@level` spec. The `@level` part is optional and defaults to the
+/// codec's lowest level.
+fn parse_spec(spec: &str) -> Result<(Codec, u8), Error> {
+    let Some((name, level)) = spec.split_once('@') else {
+        let codec = parse_codec(spec)?;
+        return Ok((codec, default_level(codec)));
+    };
+    let level = level
+        .parse::<u8>()
+        .map_err(|_| Error::Spec(format!("bad level in {spec}")))?;
+    Ok((parse_codec(name)?, level))
+}
+
+fn parse_codec(name: &str) -> Result<Codec, Error> {
+    Codec::from_name(name).ok_or_else(|| Error::Spec(format!("unknown codec: {name}")))
+}
+
+fn default_level(codec: Codec) -> u8 {
+    codec.level_range().first_level as u8
+}