From 83919bf76421fb51c21f800c7a4fb96ceba6c35a Mon Sep 17 00:00:00 2001 From: tyler Date: Mon, 25 May 2026 21:37:05 -0400 Subject: [PATCH 1/2] Switch to Tokio --- Cargo.lock | 11 +++- ruhvro/Cargo.toml | 2 +- ruhvro/benches/deserialize.rs | 79 ++++++++++++++------------ ruhvro/benches/serialize.rs | 80 ++++++++++++++------------ ruhvro/src/deserialize.rs | 102 ++++++++++++++++++++++++++++------ ruhvro/src/lib.rs | 6 ++ ruhvro/src/serialize.rs | 71 ++++++++++++++++++++--- scripts/benchmark_sweep.py | 76 +++++++++++++++++++++++++ src/lib.rs | 37 +++++++++++- 9 files changed, 362 insertions(+), 102 deletions(-) create mode 100644 scripts/benchmark_sweep.py diff --git a/Cargo.lock b/Cargo.lock index a584a36..7b49dbd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1228,7 +1228,7 @@ dependencies = [ "apache-avro", "arrow", "criterion", - "rayon", + "tokio", ] [[package]] @@ -1412,6 +1412,15 @@ dependencies = [ "serde_json", ] +[[package]] +name = "tokio" +version = "1.52.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +dependencies = [ + "pin-project-lite", +] + [[package]] name = "typenum" version = "1.17.0" diff --git a/ruhvro/Cargo.toml b/ruhvro/Cargo.toml index c5c957c..3ab852d 100644 --- a/ruhvro/Cargo.toml +++ b/ruhvro/Cargo.toml @@ -8,7 +8,7 @@ description = "Fast, multi-threaded deserialization of schema-less avro encoded repository = "https://github.com/Tyler-Sch/pyruhvro" [dependencies] -rayon = "1.10" +tokio = { version = "1", features = ["rt-multi-thread"] } apache-avro = "0.21" arrow = "58" anyhow = "1.0" diff --git a/ruhvro/benches/deserialize.rs b/ruhvro/benches/deserialize.rs index c2d26c2..b2fe99a 100644 --- a/ruhvro/benches/deserialize.rs +++ b/ruhvro/benches/deserialize.rs @@ -1,43 +1,48 @@ //! Criterion benchmarks for ruhvro's deserialize path. -//! -//! Each bench: -//! 1. Builds N avro-encoded records once (in setup). -//! 2. Measures `per_datum_deserialize` (single-threaded) and -//! `per_datum_deserialize_threaded` (8 chunks). -//! -//! Run with `cargo bench -p ruhvro --bench deserialize`. -//! HTML reports land in `target/criterion`. +//! Compares single-threaded, spawn_blocking, tokio::spawn, and rayon +//! across 1k / 10k / 100k record counts. mod common; use apache_avro::Schema as AvroSchema; use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; -use ruhvro::deserialize::{per_datum_deserialize, per_datum_deserialize_threaded}; +use ruhvro::deserialize::{ + per_datum_deserialize, + per_datum_deserialize_threaded, + per_datum_deserialize_threaded_spawn, +}; -const N: usize = 1_000; +const SIZES: &[usize] = &[1_000, 10_000]; -fn run_group(c: &mut Criterion, name: &str, parsed: AvroSchema, encoded: Vec>) { +fn run_group( + c: &mut Criterion, + schema_name: &str, + parsed: AvroSchema, + encoded: Vec>, + n: usize, +) { let refs: Vec<&[u8]> = encoded.iter().map(Vec::as_slice).collect(); + let group_name = format!("{schema_name}/{n}"); + let mut group = c.benchmark_group(&group_name); + group.throughput(Throughput::Elements(n as u64)); - let mut group = c.benchmark_group(name); - group.throughput(Throughput::Elements(N as u64)); + group.bench_function("single_threaded", |b| { + b.iter(|| black_box(per_datum_deserialize(black_box(&refs), black_box(&parsed)).unwrap())) + }); - group.bench_function("per_datum_deserialize", |b| { + group.bench_function("spawn_blocking", |b| { b.iter(|| { - let rb = per_datum_deserialize(black_box(&refs), black_box(&parsed)).unwrap(); - black_box(rb); + black_box(per_datum_deserialize_threaded( + black_box(refs.clone()), black_box(&parsed), common::NUM_CHUNKS, + ).unwrap()) }) }); - group.bench_function("per_datum_deserialize_threaded", |b| { + group.bench_function("tokio_spawn", |b| { b.iter(|| { - let rbs = per_datum_deserialize_threaded( - black_box(refs.clone()), - black_box(&parsed), - common::NUM_CHUNKS, - ) - .unwrap(); - black_box(rbs); + black_box(per_datum_deserialize_threaded_spawn( + black_box(refs.clone()), black_box(&parsed), common::NUM_CHUNKS, + ).unwrap()) }) }); @@ -45,30 +50,30 @@ fn run_group(c: &mut Criterion, name: &str, parsed: AvroSchema, encoded: Vec]) -> RecordBatch { let refs: Vec<&[u8]> = encoded.iter().map(Vec::as_slice).collect(); per_datum_deserialize(&refs, parsed).unwrap() } -fn run_group(c: &mut Criterion, name: &str, parsed: AvroSchema, batch: RecordBatch) { - let mut group = c.benchmark_group(name); - group.throughput(Throughput::Elements(N as u64)); +fn run_group(c: &mut Criterion, schema_name: &str, parsed: AvroSchema, batch: RecordBatch, n: usize) { + let group_name = format!("{schema_name}/{n}"); + let mut group = c.benchmark_group(&group_name); + group.throughput(Throughput::Elements(n as u64)); + + group.bench_function("single_threaded", |b| { + b.iter(|| { + black_box(serialize_record_batch(black_box(batch.clone()), black_box(&parsed), 1).unwrap()) + }) + }); - group.bench_function("serialize_record_batch_1chunk", |b| { + group.bench_function("spawn_blocking", |b| { b.iter(|| { - // RecordBatch holds Arc'd columns, so .clone() is cheap (refcount bumps). - let bytes = serialize_record_batch(black_box(batch.clone()), black_box(&parsed), 1) - .unwrap(); - black_box(bytes); + black_box(serialize_record_batch( + black_box(batch.clone()), black_box(&parsed), common::NUM_CHUNKS, + ).unwrap()) }) }); - group.bench_function("serialize_record_batch_8chunks", |b| { + group.bench_function("tokio_spawn", |b| { b.iter(|| { - let bytes = serialize_record_batch( - black_box(batch.clone()), - black_box(&parsed), - common::NUM_CHUNKS, - ) - .unwrap(); - black_box(bytes); + black_box(serialize_record_batch_spawn( + black_box(batch.clone()), black_box(&parsed), common::NUM_CHUNKS, + ).unwrap()) }) }); @@ -52,27 +52,35 @@ fn run_group(c: &mut Criterion, name: &str, parsed: AvroSchema, batch: RecordBat } fn bench_flat_primitives(c: &mut Criterion) { - let (parsed, encoded) = common::flat_primitives(N); - let batch = prepare_batch(&parsed, &encoded); - run_group(c, "flat_primitives", parsed, batch); + for &n in SIZES { + let (parsed, encoded) = common::flat_primitives(n); + let batch = prepare_batch(&parsed, &encoded); + run_group(c, "flat_primitives", parsed, batch, n); + } } fn bench_nullable_primitives(c: &mut Criterion) { - let (parsed, encoded) = common::nullable_primitives(N); - let batch = prepare_batch(&parsed, &encoded); - run_group(c, "nullable_primitives", parsed, batch); + for &n in SIZES { + let (parsed, encoded) = common::nullable_primitives(n); + let batch = prepare_batch(&parsed, &encoded); + run_group(c, "nullable_primitives", parsed, batch, n); + } } fn bench_nested_struct(c: &mut Criterion) { - let (parsed, encoded) = common::nested_struct(N); - let batch = prepare_batch(&parsed, &encoded); - run_group(c, "nested_struct", parsed, batch); + for &n in SIZES { + let (parsed, encoded) = common::nested_struct(n); + let batch = prepare_batch(&parsed, &encoded); + run_group(c, "nested_struct", parsed, batch, n); + } } fn bench_array_and_map(c: &mut Criterion) { - let (parsed, encoded) = common::array_and_map(N); - let batch = prepare_batch(&parsed, &encoded); - run_group(c, "array_and_map", parsed, batch); + for &n in SIZES { + let (parsed, encoded) = common::array_and_map(n); + let batch = prepare_batch(&parsed, &encoded); + run_group(c, "array_and_map", parsed, batch, n); + } } criterion_group!( diff --git a/ruhvro/src/deserialize.rs b/ruhvro/src/deserialize.rs index 1f07fb5..98d193d 100644 --- a/ruhvro/src/deserialize.rs +++ b/ruhvro/src/deserialize.rs @@ -7,7 +7,7 @@ use anyhow::{anyhow, Result}; use apache_avro::from_avro_datum; use apache_avro::Schema as AvroSchema; use arrow::array::{Array, BinaryArray, RecordBatch}; -use rayon::prelude::*; +use tokio::task; // TODO: binary array // TODO: refactor full avro deserialization to lib.rs // TODO: Durations @@ -73,26 +73,96 @@ pub fn per_datum_deserialize_threaded( slices.push(arr.slice(i * chunk_size, chunk_size)); } } - slices - .par_iter() - .map(|da| -> Result { - let chunk_refs: Vec<&[u8]> = da - .iter() - .map(|x| x.ok_or_else(|| anyhow!("Error getting sliced data"))) - .collect::>>()?; - if use_fast { - fast_decode::decode_with_arrow_schema( - &chunk_refs, - schema, - arrow_schema.as_ref().unwrap(), - ) + let schema_arc = Arc::new(schema.clone()); + crate::runtime().block_on(async { + let handles: Vec<_> = slices + .into_iter() + .map(|da| { + let schema = Arc::clone(&schema_arc); + let arrow_schema = arrow_schema.clone(); + task::spawn_blocking(move || -> Result { + let chunk_refs: Vec<&[u8]> = da + .iter() + .map(|x| x.ok_or_else(|| anyhow!("Error getting sliced data"))) + .collect::>>()?; + if use_fast { + fast_decode::decode_with_arrow_schema( + &chunk_refs, + &schema, + arrow_schema.as_ref().unwrap(), + ) + } else { + per_datum_deserialize_baseline(&chunk_refs, &schema) + } + }) + }) + .collect(); + let mut results = Vec::with_capacity(handles.len()); + for handle in handles { + results.push(handle.await.map_err(|e| anyhow!("join error: {e}"))??); + } + Ok(results) + }) +} + +/// Same as [`per_datum_deserialize_threaded`] but uses `tokio::spawn` (work-stealing +/// async pool) instead of `spawn_blocking`. CPU work runs directly on executor threads +/// with no yield points — fine for benchmarking, not for mixed async/CPU workloads. +pub fn per_datum_deserialize_threaded_spawn( + data: Vec<&[u8]>, + schema: &AvroSchema, + num_chunks: usize, +) -> Result> { + let use_fast = fast_decode::is_supported(schema); + let arrow_schema = if use_fast { + Some(Arc::new(to_arrow_schema(schema)?)) + } else { + None + }; + let arr = Arc::new(BinaryArray::from_vec(data)); + let chunk_size = arr.len() / num_chunks; + let slices: Vec<_> = (0..num_chunks) + .map(|i| { + if i == num_chunks - 1 { + arr.slice(i * chunk_size, arr.len() - (i * chunk_size)) } else { - per_datum_deserialize_baseline(&chunk_refs, schema) + arr.slice(i * chunk_size, chunk_size) } }) - .collect() + .collect(); + let schema_arc = Arc::new(schema.clone()); + crate::runtime().block_on(async { + let handles: Vec<_> = slices + .into_iter() + .map(|da| { + let schema = Arc::clone(&schema_arc); + let arrow_schema = arrow_schema.clone(); + tokio::spawn(async move { + let chunk_refs: Vec<&[u8]> = da + .iter() + .map(|x| x.ok_or_else(|| anyhow!("Error getting sliced data"))) + .collect::>>()?; + if use_fast { + fast_decode::decode_with_arrow_schema( + &chunk_refs, + &schema, + arrow_schema.as_ref().unwrap(), + ) + } else { + per_datum_deserialize_baseline(&chunk_refs, &schema) + } + }) + }) + .collect(); + let mut results = Vec::with_capacity(handles.len()); + for handle in handles { + results.push(handle.await.map_err(|e| anyhow!("join error: {e}"))??); + } + Ok(results) + }) } + #[cfg(test)] mod tests { use super::*; diff --git a/ruhvro/src/lib.rs b/ruhvro/src/lib.rs index ee8ac59..5eb309c 100644 --- a/ruhvro/src/lib.rs +++ b/ruhvro/src/lib.rs @@ -8,6 +8,12 @@ mod complex; mod fast_decode; mod fast_encode; + +use std::sync::OnceLock; +static TOKIO_RT: OnceLock = OnceLock::new(); +pub(crate) fn runtime() -> &'static tokio::runtime::Runtime { + TOKIO_RT.get_or_init(|| tokio::runtime::Runtime::new().expect("failed to build tokio runtime")) +} /// Converts Avro to Arrow /// Decode Avro data returning an Arrow Record Batch. /// ## Example diff --git a/ruhvro/src/serialize.rs b/ruhvro/src/serialize.rs index 9b00fb7..8ab61b1 100644 --- a/ruhvro/src/serialize.rs +++ b/ruhvro/src/serialize.rs @@ -2,10 +2,10 @@ use apache_avro::Schema; use arrow::array::{ Array, ArrayRef, GenericBinaryArray, RecordBatch, StructArray, }; -use rayon::prelude::*; +use tokio::task; use std::sync::Arc; use crate::serialization_containers; -use anyhow::Result; +use anyhow::{anyhow, Result}; // TODO: Should be checks to make sure avro and arrow schema match // TODO: need to figure out names. Should serialize match by name or position? // TODO: Should it include namespace in matching? @@ -43,18 +43,71 @@ pub fn serialize_record_batch( } }) .collect(); - slices - .par_iter() - .map(|x| { - if use_fast { - crate::fast_encode::serialize_chunk(schema, x) + let schema_arc = Arc::new(schema.clone()); + crate::runtime().block_on(async { + let handles: Vec<_> = slices + .into_iter() + .map(|x| { + let schema = Arc::clone(&schema_arc); + task::spawn_blocking(move || { + if use_fast { + crate::fast_encode::serialize_chunk(&schema, &x) + } else { + serialization_containers::serialize(&schema, &x) + } + }) + }) + .collect(); + let mut results = Vec::with_capacity(handles.len()); + for handle in handles { + results.push(handle.await.map_err(|e| anyhow!("join error: {e}"))??); + } + Ok(results) + }) +} + +/// Same as [`serialize_record_batch`] but uses `tokio::spawn` (work-stealing async pool). +pub fn serialize_record_batch_spawn( + rb: RecordBatch, + schema: &Schema, + num_chunks: usize, +) -> Result>> { + let use_fast = crate::fast_encode::is_supported(schema); + let schema_arc = Arc::new(schema.clone()); + let struct_arry: ArrayRef = Arc::::new(rb.into()); + let chunk_size = struct_arry.len() / num_chunks; + let slices: Vec<_> = (0..num_chunks) + .map(|i| { + if i == num_chunks - 1 { + struct_arry.slice(i * chunk_size, struct_arry.len() - (i * chunk_size)) } else { - serialization_containers::serialize(schema, x) + struct_arry.slice(i * chunk_size, chunk_size) } }) - .collect() + .collect(); + crate::runtime().block_on(async { + let handles: Vec<_> = slices + .into_iter() + .map(|x| { + let schema = Arc::clone(&schema_arc); + tokio::spawn(async move { + if use_fast { + crate::fast_encode::serialize_chunk(&schema, &x) + } else { + serialization_containers::serialize(&schema, &x) + } + }) + }) + .collect(); + let mut results = Vec::with_capacity(handles.len()); + for handle in handles { + results.push(handle.await.map_err(|e| anyhow!("join error: {e}"))??); + } + Ok(results) + }) } + #[cfg(test)] mod test { use crate::serialize::serialize_record_batch; diff --git a/scripts/benchmark_sweep.py b/scripts/benchmark_sweep.py new file mode 100644 index 0000000..c0481eb --- /dev/null +++ b/scripts/benchmark_sweep.py @@ -0,0 +1,76 @@ +"""Benchmark pyruhvro vs fastavro across dataset sizes and chunk counts.""" +import io +import json +import timeit + +import fastavro +from pyruhvro import deserialize_array_threaded, serialize_record_batch + +import generate_avro + +RECORD_COUNTS = [500, 5_000, 50_000] +CHUNK_COUNTS = [1, 2, 4, 8, 16] +TIMEIT_REPS = 3 # best-of-N repeats per measurement +MIN_SECONDS = 1.5 # timeit runs enough loops to fill at least this long + + +def measure(stmt, setup_vars, seconds=MIN_SECONDS, reps=TIMEIT_REPS): + """Return best ms-per-loop over `reps` rounds of auto-scaled timeit.""" + t = timeit.Timer(stmt=stmt, globals=setup_vars) + n, _ = t.autorange() + n = max(n, 3) + times = t.repeat(reps, n) + return min(times) / n * 1000 # ms + + +def fastavro_serialize(records, parsed_schema): + out = [] + for r in records: + buf = io.BytesIO() + fastavro.schemaless_writer(buf, parsed_schema, r) + out.append(buf.getvalue()) + return out + + +def fastavro_deserialize(serialized, parsed_schema): + return [fastavro.schemaless_reader(io.BytesIO(b), parsed_schema) for b in serialized] + + +print(f"\n{'Dataset':>8} {'Chunks':>6} {'py-ser ms':>10} {'py-de ms':>10} {'fa-ser ms':>10} {'fa-de ms':>10} {'ser-x':>6} {'de-x':>5}") +print("-" * 88) + +for n_rec in RECORD_COUNTS: + records = generate_avro.generate_records(n_rec) + serialized = generate_avro.get_serialized_records(records) + parsed_schema = generate_avro.parsed_schema + schema_str = generate_avro.schema_string + + # fastavro baseline (no chunks, run once per dataset size) + g = dict(records=records, serialized=serialized, parsed_schema=parsed_schema, + fastavro_serialize=fastavro_serialize, fastavro_deserialize=fastavro_deserialize) + fa_ser = measure("fastavro_serialize(records, parsed_schema)", g) + fa_de = measure("fastavro_deserialize(serialized, parsed_schema)", g) + + first_chunk = True + for n_chunks in CHUNK_COUNTS: + g2 = dict(serialized=serialized, schema_str=schema_str, + serialize_record_batch=serialize_record_batch, + deserialize_array_threaded=deserialize_array_threaded, + n_chunks=n_chunks) + + # pre-deserialize so serialize bench has record_batches ready + record_batches = deserialize_array_threaded(serialized, schema_str, n_chunks) + g2["record_batches"] = record_batches + + py_ser = measure("[ serialize_record_batch(rb, schema_str, n_chunks) for rb in record_batches ]", g2) + py_de = measure("deserialize_array_threaded(serialized, schema_str, n_chunks)", g2) + + fa_ser_col = f"{fa_ser:>10.2f}" if first_chunk else f"{'':>10}" + fa_de_col = f"{fa_de:>10.2f}" if first_chunk else f"{'':>10}" + first_chunk = False + + ser_x = fa_ser / py_ser + de_x = fa_de / py_de + print(f"{n_rec:>8,} {n_chunks:>6} {py_ser:>10.2f} {py_de:>10.2f} {fa_ser_col} {fa_de_col} {ser_x:>5.1f}x {de_x:>4.1f}x") + + print() diff --git a/src/lib.rs b/src/lib.rs index 988a0a9..32da027 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,8 +7,7 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::pybacked::PyBackedBytes; use pyo3::types::PyList; -use ruhvro::deserialize; -use ruhvro::serialize; +use ruhvro::{deserialize, serialize}; fn to_py_err(e: E) -> PyErr { PyValueError::new_err(e.to_string()) @@ -60,11 +59,45 @@ fn serialize_record_batch( .collect()) } +#[pyfunction] +fn deserialize_array_threaded_spawn( + list: &Bound<'_, PyList>, + schema: &str, + num_chunks: usize, +) -> PyResult>> { + let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; + let owned = extract_bytes_list(list)?; + let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); + let record_batches = + deserialize::per_datum_deserialize_threaded_spawn(borrow_list, &parsed_schema, num_chunks) + .map_err(to_py_err)?; + Ok(record_batches.into_iter().map(PyArrowType).collect()) +} + +#[pyfunction] +fn serialize_record_batch_spawn( + data: PyArrowType, + schema: &str, + num_chunks: usize, +) -> PyResult>> { + let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; + let serialized = + serialize::serialize_record_batch_spawn(data.0, &parsed_schema, num_chunks) + .map_err(to_py_err)?; + Ok(serialized + .into_iter() + .map(|x| PyArrowType(x.into_data())) + .collect()) +} + + /// A Python module implemented in Rust. #[pymodule] fn pyruhvro(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(deserialize_array, m)?)?; m.add_function(wrap_pyfunction!(deserialize_array_threaded, m)?)?; m.add_function(wrap_pyfunction!(serialize_record_batch, m)?)?; + m.add_function(wrap_pyfunction!(deserialize_array_threaded_spawn, m)?)?; + m.add_function(wrap_pyfunction!(serialize_record_batch_spawn, m)?)?; Ok(()) } From af82d42fc984c67ba793645c3e880061362e9bea Mon Sep 17 00:00:00 2001 From: tyler Date: Mon, 25 May 2026 22:51:18 -0400 Subject: [PATCH 2/2] Fix schema issue --- CLAUDE.md | 11 ++-- ruhvro/benches/deserialize.rs | 7 ++- ruhvro/benches/serialize.rs | 11 +++- ruhvro/examples/prof_decode.rs | 3 +- ruhvro/src/deserialize.rs | 71 ++++++++++++----------- ruhvro/src/fast_encode.rs | 4 +- ruhvro/src/lib.rs | 13 +++-- ruhvro/src/serialize.rs | 85 +++++++++++++-------------- src/lib.rs | 101 +++++++++++++++++++++++++-------- 9 files changed, 188 insertions(+), 118 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 2c84520..c023c31 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -7,7 +7,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co Cargo workspace with two crates: - `ruhvro/` — the core Rust library (published as `ruhvro` on crates.io). Pure Rust API for serializing/deserializing schemaless Avro to/from Arrow `RecordBatch`es. Has no Python deps. -- `src/lib.rs` (top-level) — the `pyruhvro` PyO3 extension module that wraps `ruhvro` and exposes it to Python via maturin. Re-exports three functions: `deserialize_array`, `deserialize_array_threaded`, `serialize_record_batch`. +- `src/lib.rs` (top-level) — the `pyruhvro` PyO3 extension module that wraps `ruhvro` and exposes it to Python via maturin. Exposes `deserialize_array`, `deserialize_array_threaded`, `serialize_record_batch`, plus `_spawn` variants that use `tokio::spawn` instead of `spawn_blocking`. Also maintains a `String -> Arc` cache so repeat calls don't re-parse the schema JSON, and releases the GIL via `Python::detach` around every Rust call so multiple Python threads can run concurrently. Keep Python-facing concerns in the top-level crate; keep the Avro↔Arrow logic in `ruhvro/`. The PyO3 wrappers should stay thin — convert PyArrow types, call into `ruhvro`, return PyArrow types. @@ -42,15 +42,18 @@ The pipeline is **schemaless Avro bytes ⇄ Arrow `RecordBatch`**, driven by a p Key modules in `ruhvro/src/`: -- `deserialize.rs` — public entry points `parse_schema`, `per_datum_deserialize` (single-threaded), `per_datum_deserialize_threaded` (rayon, splits input into `num_chunks` slices and returns one `RecordBatch` per chunk). -- `serialize.rs` — public entry point `serialize_record_batch`. Converts the `RecordBatch` into a `StructArray`, slices it into `num_chunks`, and serializes each slice in parallel via rayon into a `GenericBinaryArray` of Avro datums. +- `deserialize.rs` — public entry points `parse_schema`, `per_datum_deserialize` (single-threaded), `per_datum_deserialize_threaded` (tokio `spawn_blocking`, splits input into `num_chunks` slices and returns one `RecordBatch` per chunk), plus `_spawn` variant on the work-stealing async pool. Threaded variants take an `Arc` so callers can share one parsed schema across many calls without re-cloning. +- `serialize.rs` — public entry point `serialize_record_batch` (same `Arc` convention). Converts the `RecordBatch` into a `StructArray`, slices it into `num_chunks`, and serializes each slice in parallel via tokio `spawn_blocking` into a `GenericBinaryArray` of Avro datums. +- `fast_decode.rs` / `fast_encode.rs` — schema-walking decoder/encoder that bypasses the `apache_avro::Value` tree and writes straight into Arrow builders. Gated by `is_supported(schema)`; falls back to the `Value`-based path for schemas containing types outside the supported subset. - `schema_translate.rs` — converts an `apache_avro::Schema` into an `arrow::datatypes::Schema`. This is the source of truth for type mapping (e.g. nullable-union → nullable Arrow field, multi-variant unions → Arrow `Union`, Avro `map` → Arrow `Map`, logical types → Arrow temporal types). - `complex.rs` — `AvroToArrowBuilder` and its `Struct`/`List`/`Union`/`Map`/`Primitive` variants. This is the *deserialize* side: walks Avro `Value`s into Arrow builders. The `add_val!` macro and `get_val_from_possible_union` helper handle the common "value might be wrapped in a union" case. - `serialization_containers.rs` — the *serialize* side: `ArrayContainers` walks Arrow arrays column-wise and re-emits `apache_avro::types::Value`s, then `to_avro_datum` encodes each row. ### Threading model -Both threaded paths use rayon and require the caller to pass `num_chunks` explicitly — the library does not infer it from `rayon::current_num_threads()`. The Python wrappers release the GIL implicitly via PyO3 while inside `ruhvro` calls, which is the whole point of the project per the README. +A single global tokio multi-thread runtime (`OnceLock` in `ruhvro/src/lib.rs`) services all parallel work — created on first call, alive for the process lifetime. Worker threads are parked when idle, so the runtime costs nothing when unused. + +Both threaded paths require the caller to pass `num_chunks` explicitly — the library does not infer it. `num_chunks` is clamped to `[1, max(rows, 1)]` so `0` doesn't panic and overshooting the row count doesn't spawn empty tasks. The Python wrappers explicitly release the GIL with `py.detach(...)` around every Rust call, so multiple Python threads can call into pyruhvro concurrently and benefit from the internal parallelism. ### Avro union handling diff --git a/ruhvro/benches/deserialize.rs b/ruhvro/benches/deserialize.rs index b2fe99a..a8c40c0 100644 --- a/ruhvro/benches/deserialize.rs +++ b/ruhvro/benches/deserialize.rs @@ -4,6 +4,8 @@ mod common; +use std::sync::Arc; + use apache_avro::Schema as AvroSchema; use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; use ruhvro::deserialize::{ @@ -22,6 +24,7 @@ fn run_group( n: usize, ) { let refs: Vec<&[u8]> = encoded.iter().map(Vec::as_slice).collect(); + let schema_arc = Arc::new(parsed.clone()); let group_name = format!("{schema_name}/{n}"); let mut group = c.benchmark_group(&group_name); group.throughput(Throughput::Elements(n as u64)); @@ -33,7 +36,7 @@ fn run_group( group.bench_function("spawn_blocking", |b| { b.iter(|| { black_box(per_datum_deserialize_threaded( - black_box(refs.clone()), black_box(&parsed), common::NUM_CHUNKS, + black_box(refs.clone()), black_box(Arc::clone(&schema_arc)), common::NUM_CHUNKS, ).unwrap()) }) }); @@ -41,7 +44,7 @@ fn run_group( group.bench_function("tokio_spawn", |b| { b.iter(|| { black_box(per_datum_deserialize_threaded_spawn( - black_box(refs.clone()), black_box(&parsed), common::NUM_CHUNKS, + black_box(refs.clone()), black_box(Arc::clone(&schema_arc)), common::NUM_CHUNKS, ).unwrap()) }) }); diff --git a/ruhvro/benches/serialize.rs b/ruhvro/benches/serialize.rs index 78f493b..ceb1ba1 100644 --- a/ruhvro/benches/serialize.rs +++ b/ruhvro/benches/serialize.rs @@ -5,6 +5,8 @@ mod common; +use std::sync::Arc; + use apache_avro::Schema as AvroSchema; use arrow::array::RecordBatch; use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; @@ -22,20 +24,23 @@ fn prepare_batch(parsed: &AvroSchema, encoded: &[Vec]) -> RecordBatch { } fn run_group(c: &mut Criterion, schema_name: &str, parsed: AvroSchema, batch: RecordBatch, n: usize) { + let schema_arc = Arc::new(parsed); let group_name = format!("{schema_name}/{n}"); let mut group = c.benchmark_group(&group_name); group.throughput(Throughput::Elements(n as u64)); group.bench_function("single_threaded", |b| { b.iter(|| { - black_box(serialize_record_batch(black_box(batch.clone()), black_box(&parsed), 1).unwrap()) + black_box(serialize_record_batch( + black_box(batch.clone()), black_box(Arc::clone(&schema_arc)), 1, + ).unwrap()) }) }); group.bench_function("spawn_blocking", |b| { b.iter(|| { black_box(serialize_record_batch( - black_box(batch.clone()), black_box(&parsed), common::NUM_CHUNKS, + black_box(batch.clone()), black_box(Arc::clone(&schema_arc)), common::NUM_CHUNKS, ).unwrap()) }) }); @@ -43,7 +48,7 @@ fn run_group(c: &mut Criterion, schema_name: &str, parsed: AvroSchema, batch: Re group.bench_function("tokio_spawn", |b| { b.iter(|| { black_box(serialize_record_batch_spawn( - black_box(batch.clone()), black_box(&parsed), common::NUM_CHUNKS, + black_box(batch.clone()), black_box(Arc::clone(&schema_arc)), common::NUM_CHUNKS, ).unwrap()) }) }); diff --git a/ruhvro/examples/prof_decode.rs b/ruhvro/examples/prof_decode.rs index 5754fe9..c93ca05 100644 --- a/ruhvro/examples/prof_decode.rs +++ b/ruhvro/examples/prof_decode.rs @@ -40,6 +40,7 @@ fn build_record(i: usize, sch: &AvroSchema) -> Value { fn main() { let parsed = parse_schema(SCHEMA).unwrap(); + let schema_arc = std::sync::Arc::new(parsed.clone()); let n = 1_000; let encoded: Vec> = (0..n) .map(|i| to_avro_datum(&parsed, build_record(i, &parsed)).unwrap()) @@ -49,7 +50,7 @@ fn main() { let mut total_rows: usize = 0; for iter in 0..10_000 { let refs: Vec<&[u8]> = encoded.iter().map(Vec::as_slice).collect(); - let rbs = per_datum_deserialize_threaded(refs, &parsed, 8).unwrap(); + let rbs = per_datum_deserialize_threaded(refs, std::sync::Arc::clone(&schema_arc), 8).unwrap(); total_rows += rbs.iter().map(|rb| rb.num_rows()).sum::(); if iter % 1000 == 0 { eprintln!(" iter {iter}"); diff --git a/ruhvro/src/deserialize.rs b/ruhvro/src/deserialize.rs index 98d193d..a704a6b 100644 --- a/ruhvro/src/deserialize.rs +++ b/ruhvro/src/deserialize.rs @@ -47,38 +47,53 @@ pub(crate) fn per_datum_deserialize_baseline( Ok(sa.into()) } +/// Clamp `num_chunks` to the range `[1, max(data_len, 1)]`. Callers that pass +/// `0` get a single chunk; callers asking for more chunks than rows get one +/// chunk per row so we don't spawn empty tasks. +fn clamp_chunks(num_chunks: usize, data_len: usize) -> usize { + num_chunks.max(1).min(data_len.max(1)) +} + +fn build_slices(arr: &Arc, num_chunks: usize) -> Vec { + let chunk_size = arr.len() / num_chunks; + (0..num_chunks) + .map(|i| { + if i == num_chunks - 1 { + arr.slice(i * chunk_size, arr.len() - (i * chunk_size)) + } else { + arr.slice(i * chunk_size, chunk_size) + } + }) + .collect() +} + /// Deserializes vector of serialized avro messages with many threads. /// Each chunk dispatches to the [`fast_decode`] path when the schema is supported. +/// +/// The caller passes the parsed schema in an `Arc` so it can be cheaply shared +/// across spawned tasks. Reusing the same `Arc` across calls avoids re-cloning +/// the schema on every invocation. pub fn per_datum_deserialize_threaded( data: Vec<&[u8]>, - schema: &AvroSchema, + schema: Arc, num_chunks: usize, ) -> Result> { - let use_fast = fast_decode::is_supported(schema); + let num_chunks = clamp_chunks(num_chunks, data.len()); + let use_fast = fast_decode::is_supported(&schema); // Compute the Arrow schema once and share it across all chunks — avoids // an `to_arrow_schema` walk per chunk on the fast path. let arrow_schema = if use_fast { - Some(Arc::new(to_arrow_schema(schema)?)) + Some(Arc::new(to_arrow_schema(&schema)?)) } else { None }; let arr = Arc::new(BinaryArray::from_vec(data)); - let mut slices = vec![]; - let cores = num_chunks; - let chunk_size = arr.len() / cores; - for i in 0..cores { - if i == cores - 1 { - slices.push(arr.slice(i * chunk_size, arr.len() - (i * chunk_size))); - } else { - slices.push(arr.slice(i * chunk_size, chunk_size)); - } - } - let schema_arc = Arc::new(schema.clone()); + let slices = build_slices(&arr, num_chunks); crate::runtime().block_on(async { let handles: Vec<_> = slices .into_iter() .map(|da| { - let schema = Arc::clone(&schema_arc); + let schema = Arc::clone(&schema); let arrow_schema = arrow_schema.clone(); task::spawn_blocking(move || -> Result { let chunk_refs: Vec<&[u8]> = da @@ -106,36 +121,28 @@ pub fn per_datum_deserialize_threaded( } /// Same as [`per_datum_deserialize_threaded`] but uses `tokio::spawn` (work-stealing -/// async pool) instead of `spawn_blocking`. CPU work runs directly on executor threads -/// with no yield points — fine for benchmarking, not for mixed async/CPU workloads. +/// async pool) instead of `spawn_blocking`. CPU work runs directly on executor +/// threads with no yield points; safe to use here because the global runtime +/// only services ruhvro tasks. pub fn per_datum_deserialize_threaded_spawn( data: Vec<&[u8]>, - schema: &AvroSchema, + schema: Arc, num_chunks: usize, ) -> Result> { - let use_fast = fast_decode::is_supported(schema); + let num_chunks = clamp_chunks(num_chunks, data.len()); + let use_fast = fast_decode::is_supported(&schema); let arrow_schema = if use_fast { - Some(Arc::new(to_arrow_schema(schema)?)) + Some(Arc::new(to_arrow_schema(&schema)?)) } else { None }; let arr = Arc::new(BinaryArray::from_vec(data)); - let chunk_size = arr.len() / num_chunks; - let slices: Vec<_> = (0..num_chunks) - .map(|i| { - if i == num_chunks - 1 { - arr.slice(i * chunk_size, arr.len() - (i * chunk_size)) - } else { - arr.slice(i * chunk_size, chunk_size) - } - }) - .collect(); - let schema_arc = Arc::new(schema.clone()); + let slices = build_slices(&arr, num_chunks); crate::runtime().block_on(async { let handles: Vec<_> = slices .into_iter() .map(|da| { - let schema = Arc::clone(&schema_arc); + let schema = Arc::clone(&schema); let arrow_schema = arrow_schema.clone(); tokio::spawn(async move { let chunk_refs: Vec<&[u8]> = da diff --git a/ruhvro/src/fast_encode.rs b/ruhvro/src/fast_encode.rs index eb4b001..7acae2c 100644 --- a/ruhvro/src/fast_encode.rs +++ b/ruhvro/src/fast_encode.rs @@ -624,7 +624,7 @@ mod tests { // Decode via fast path, then serialize via fast path. let rb: RecordBatch = crate::fast_decode::decode(&refs, &s).unwrap(); - let bytes_chunks = serialize_record_batch(rb.clone(), &s, 1).unwrap(); + let bytes_chunks = serialize_record_batch(rb.clone(), std::sync::Arc::new(s.clone()), 1).unwrap(); assert_eq!(bytes_chunks.len(), 1); let chunk = &bytes_chunks[0]; @@ -832,7 +832,7 @@ mod tests { let refs: Vec<&[u8]> = owned.iter().map(Vec::as_slice).collect(); let rb: RecordBatch = crate::fast_decode::decode(&refs, &s).unwrap(); - let bytes_chunks = serialize_record_batch(rb.clone(), &s, 1).unwrap(); + let bytes_chunks = serialize_record_batch(rb.clone(), std::sync::Arc::new(s.clone()), 1).unwrap(); let chunk = &bytes_chunks[0]; let round_refs: Vec<&[u8]> = (0..chunk.len()).map(|i| chunk.value(i)).collect(); let round = crate::fast_decode::decode(&round_refs, &s).unwrap(); diff --git a/ruhvro/src/lib.rs b/ruhvro/src/lib.rs index 5eb309c..bb3ce70 100644 --- a/ruhvro/src/lib.rs +++ b/ruhvro/src/lib.rs @@ -43,8 +43,10 @@ pub(crate) fn runtime() -> &'static tokio::runtime::Runtime { /// let deserialized = ruhvro::deserialize::per_datum_deserialize(&vec![&serialized[..]], &parsed_schema).unwrap(); /// println!("{:?}", deserialized); /// -/// // serialize the record batch -/// let serialized = ruhvro::serialize::serialize_record_batch(deserialized, &parsed_schema, 1).unwrap(); +/// // serialize the record batch — threaded functions take an `Arc` +/// // so the parsed schema can be reused cheaply across calls. +/// let schema_arc = std::sync::Arc::new(parsed_schema.clone()); +/// let serialized = ruhvro::serialize::serialize_record_batch(deserialized, schema_arc, 1).unwrap(); /// println!("{:?}", serialized); /// /// @@ -169,8 +171,11 @@ mod tests { let newv = vec![&encoded[..], &encoded2[..], &encoded3[..]]; let result = per_datum_deserialize(&newv, &parsed_schema).unwrap(); - let serialized = crate::serialize::serialize_record_batch(result, &parsed_schema, 1).unwrap(); - + let _serialized = crate::serialize::serialize_record_batch( + result, + std::sync::Arc::new(parsed_schema.clone()), + 1, + ).unwrap(); } pub fn decode_hex(s: &str) -> Vec { diff --git a/ruhvro/src/serialize.rs b/ruhvro/src/serialize.rs index 8ab61b1..49670ef 100644 --- a/ruhvro/src/serialize.rs +++ b/ruhvro/src/serialize.rs @@ -12,43 +12,43 @@ use anyhow::{anyhow, Result}; // TODO: Add check for sparse union types // TODO: remove any unwraps and check results/errors -/// Serializes a `RecordBatch` into a vector of `GenericBinaryArray`. -/// -/// This function takes a `RecordBatch` and a schema as input and serializes the data -/// in the `RecordBatch` into a vector of `GenericBinaryArray`. Each `GenericBinaryArray` -/// represents a chunk of the serialized data. -/// -/// # Arguments -/// -/// * `rb` - The `RecordBatch` to be serialized. -/// * `schema` - The schema of the `RecordBatch`. -/// -/// # Returns +fn clamp_chunks(num_chunks: usize, data_len: usize) -> usize { + num_chunks.max(1).min(data_len.max(1)) +} + +fn slice_struct(arr: &ArrayRef, num_chunks: usize) -> Vec { + let chunk_size = arr.len() / num_chunks; + (0..num_chunks) + .map(|i| { + if i == num_chunks - 1 { + arr.slice(i * chunk_size, arr.len() - (i * chunk_size)) + } else { + arr.slice(i * chunk_size, chunk_size) + } + }) + .collect() +} + +/// Serializes a `RecordBatch` into a vector of `GenericBinaryArray`, +/// one per chunk. /// -/// A vector of `GenericBinaryArray` representing the serialized data. +/// The caller passes the schema in an `Arc` so it can be cheaply shared across +/// spawned tasks. `num_chunks` is clamped to `[1, rows]` to avoid spawning +/// empty tasks when callers ask for more chunks than rows. pub fn serialize_record_batch( rb: RecordBatch, - schema: &Schema, + schema: Arc, num_chunks: usize, ) -> Result>> { - let use_fast = crate::fast_encode::is_supported(schema); + let use_fast = crate::fast_encode::is_supported(&schema); let struct_arry: ArrayRef = Arc::::new(rb.into()); - let chunk_size = struct_arry.len() / num_chunks; - let slices: Vec<_> = (0..num_chunks) - .map(|i| { - if i == num_chunks - 1 { - struct_arry.slice(i * chunk_size, struct_arry.len() - (i * chunk_size)) - } else { - struct_arry.slice(i * chunk_size, chunk_size) - } - }) - .collect(); - let schema_arc = Arc::new(schema.clone()); + let num_chunks = clamp_chunks(num_chunks, struct_arry.len()); + let slices = slice_struct(&struct_arry, num_chunks); crate::runtime().block_on(async { let handles: Vec<_> = slices .into_iter() .map(|x| { - let schema = Arc::clone(&schema_arc); + let schema = Arc::clone(&schema); task::spawn_blocking(move || { if use_fast { crate::fast_encode::serialize_chunk(&schema, &x) @@ -69,27 +69,18 @@ pub fn serialize_record_batch( /// Same as [`serialize_record_batch`] but uses `tokio::spawn` (work-stealing async pool). pub fn serialize_record_batch_spawn( rb: RecordBatch, - schema: &Schema, + schema: Arc, num_chunks: usize, ) -> Result>> { - let use_fast = crate::fast_encode::is_supported(schema); - let schema_arc = Arc::new(schema.clone()); + let use_fast = crate::fast_encode::is_supported(&schema); let struct_arry: ArrayRef = Arc::::new(rb.into()); - let chunk_size = struct_arry.len() / num_chunks; - let slices: Vec<_> = (0..num_chunks) - .map(|i| { - if i == num_chunks - 1 { - struct_arry.slice(i * chunk_size, struct_arry.len() - (i * chunk_size)) - } else { - struct_arry.slice(i * chunk_size, chunk_size) - } - }) - .collect(); + let num_chunks = clamp_chunks(num_chunks, struct_arry.len()); + let slices = slice_struct(&struct_arry, num_chunks); crate::runtime().block_on(async { let handles: Vec<_> = slices .into_iter() .map(|x| { - let schema = Arc::clone(&schema_arc); + let schema = Arc::clone(&schema); tokio::spawn(async move { if use_fast { crate::fast_encode::serialize_chunk(&schema, &x) @@ -203,7 +194,7 @@ mod test { let schema = Schema::parse_str(avro_schema).unwrap(); // let struct_arr_ref: ArrayRef = Arc::new(struct_arr); let struct_arr_ref: RecordBatch = struct_arr.into(); - let r = serialize_record_batch(struct_arr_ref.clone(), &schema, 1).unwrap(); + let r = serialize_record_batch(struct_arr_ref.clone(), Arc::new(schema.clone()), 1).unwrap(); let ra = r .iter() .map(|x| x.iter().map(|j| j.unwrap()).collect::>()) @@ -261,7 +252,7 @@ mod test { ); let arr: RecordBatch = outer_struct_arr.into(); let parsed_schmea = Schema::parse_str(schmea).unwrap(); - let arr_cont = serialize_record_batch(arr.clone(), &parsed_schmea, 1).unwrap(); + let arr_cont = serialize_record_batch(arr.clone(), Arc::new(parsed_schmea.clone()), 1).unwrap(); let ra = arr_cont .iter() .map(|x| x.iter().map(|j| j.unwrap()).collect::>()) @@ -308,7 +299,7 @@ mod test { let schema = Schema::parse_str(schema).unwrap(); let record_arr_ref: RecordBatch = record_arr.into(); println!("{:?}", record_arr_ref); - let r = serialize_record_batch(record_arr_ref.clone(), &schema, 1).unwrap(); + let r = serialize_record_batch(record_arr_ref.clone(), Arc::new(schema.clone()), 1).unwrap(); let ra = r .iter() .map(|x| x.iter().map(|j| j.unwrap()).collect::>()) @@ -398,8 +389,8 @@ mod test { ("a", a_arr.clone()), ]).unwrap(); - let in_order_bytes = serialize_record_batch(in_order, &parsed, 1).unwrap(); - let reversed_bytes = serialize_record_batch(reversed, &parsed, 1).unwrap(); + let in_order_bytes = serialize_record_batch(in_order, Arc::new(parsed.clone()), 1).unwrap(); + let reversed_bytes = serialize_record_batch(reversed, Arc::new(parsed.clone()), 1).unwrap(); assert_eq!(in_order_bytes.len(), reversed_bytes.len()); for (a, b) in in_order_bytes.iter().zip(reversed_bytes.iter()) { @@ -425,7 +416,7 @@ mod test { let a_arr: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); let batch = RecordBatch::try_from_iter(vec![("a", a_arr)]).unwrap(); - let err = serialize_record_batch(batch, &parsed, 1).unwrap_err(); + let err = serialize_record_batch(batch, Arc::new(parsed.clone()), 1).unwrap_err(); let msg = format!("{err}"); assert!(msg.contains("missing column 'b'"), "unexpected error: {msg}"); } diff --git a/src/lib.rs b/src/lib.rs index 32da027..13d6c4a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,19 @@ //! Python extensions for transforming a vector of avro encoded binary data to an -//! apache arrow record batch +//! apache arrow record batch. //! +//! Two cross-cutting concerns live here: +//! 1. **Schema cache** — `ruhvro`'s threaded API takes `Arc` so the +//! parsed schema is shared (not cloned) across spawned tasks. We cache +//! parsed schemas keyed by their source string so Python callers don't +//! re-parse JSON on every call. +//! 2. **GIL release** — every Python entry point releases the GIL around the +//! Rust work so multiple Python threads can call into pyruhvro +//! concurrently and benefit from ruhvro's internal parallelism. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, OnceLock}; + +use apache_avro::Schema as AvroSchema; use arrow::array::{Array, ArrayData, RecordBatch}; use arrow::pyarrow::PyArrowType; use pyo3::exceptions::PyValueError; @@ -19,39 +32,72 @@ fn extract_bytes_list(list: &Bound<'_, PyList>) -> PyResult> .collect() } +/// Schema cache, keyed by the raw schema string. Avoids reparsing on every +/// call — schema parsing dominates small-payload latency for repeated calls +/// with the same schema. Unbounded by design: real workloads use a handful +/// of distinct schemas per process. +fn schema_cache() -> &'static Mutex>> { + static CACHE: OnceLock>>> = OnceLock::new(); + CACHE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn get_or_parse_schema(schema: &str) -> PyResult> { + { + let cache = schema_cache().lock().expect("schema cache poisoned"); + if let Some(s) = cache.get(schema) { + return Ok(Arc::clone(s)); + } + } + let parsed = Arc::new(deserialize::parse_schema(schema).map_err(to_py_err)?); + let mut cache = schema_cache().lock().expect("schema cache poisoned"); + Ok(Arc::clone(cache.entry(schema.to_string()).or_insert(parsed))) +} + #[pyfunction] -fn deserialize_array(list: &Bound<'_, PyList>, schema: &str) -> PyResult> { - let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; +fn deserialize_array( + py: Python<'_>, + list: &Bound<'_, PyList>, + schema: &str, +) -> PyResult> { + let parsed_schema = get_or_parse_schema(schema)?; let owned = extract_bytes_list(list)?; - let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); - let record_batch = - deserialize::per_datum_deserialize(&borrow_list, &parsed_schema).map_err(to_py_err)?; + let record_batch = py + .detach(move || { + let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); + deserialize::per_datum_deserialize(&borrow_list, &parsed_schema) + }) + .map_err(to_py_err)?; Ok(PyArrowType(record_batch)) } #[pyfunction] fn deserialize_array_threaded( + py: Python<'_>, list: &Bound<'_, PyList>, schema: &str, num_chunks: usize, ) -> PyResult>> { - let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; + let parsed_schema = get_or_parse_schema(schema)?; let owned = extract_bytes_list(list)?; - let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); - let record_batches = - deserialize::per_datum_deserialize_threaded(borrow_list, &parsed_schema, num_chunks) - .map_err(to_py_err)?; + let record_batches = py + .detach(move || { + let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); + deserialize::per_datum_deserialize_threaded(borrow_list, parsed_schema, num_chunks) + }) + .map_err(to_py_err)?; Ok(record_batches.into_iter().map(PyArrowType).collect()) } #[pyfunction] fn serialize_record_batch( + py: Python<'_>, data: PyArrowType, schema: &str, num_chunks: usize, ) -> PyResult>> { - let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; - let serialized = serialize::serialize_record_batch(data.0, &parsed_schema, num_chunks) + let parsed_schema = get_or_parse_schema(schema)?; + let serialized = py + .detach(move || serialize::serialize_record_batch(data.0, parsed_schema, num_chunks)) .map_err(to_py_err)?; Ok(serialized .into_iter() @@ -61,36 +107,45 @@ fn serialize_record_batch( #[pyfunction] fn deserialize_array_threaded_spawn( + py: Python<'_>, list: &Bound<'_, PyList>, schema: &str, num_chunks: usize, ) -> PyResult>> { - let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; + let parsed_schema = get_or_parse_schema(schema)?; let owned = extract_bytes_list(list)?; - let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); - let record_batches = - deserialize::per_datum_deserialize_threaded_spawn(borrow_list, &parsed_schema, num_chunks) - .map_err(to_py_err)?; + let record_batches = py + .detach(move || { + let borrow_list: Vec<&[u8]> = owned.iter().map(|b| &b[..]).collect(); + deserialize::per_datum_deserialize_threaded_spawn( + borrow_list, + parsed_schema, + num_chunks, + ) + }) + .map_err(to_py_err)?; Ok(record_batches.into_iter().map(PyArrowType).collect()) } #[pyfunction] fn serialize_record_batch_spawn( + py: Python<'_>, data: PyArrowType, schema: &str, num_chunks: usize, ) -> PyResult>> { - let parsed_schema = deserialize::parse_schema(schema).map_err(to_py_err)?; - let serialized = - serialize::serialize_record_batch_spawn(data.0, &parsed_schema, num_chunks) - .map_err(to_py_err)?; + let parsed_schema = get_or_parse_schema(schema)?; + let serialized = py + .detach(move || { + serialize::serialize_record_batch_spawn(data.0, parsed_schema, num_chunks) + }) + .map_err(to_py_err)?; Ok(serialized .into_iter() .map(|x| PyArrowType(x.into_data())) .collect()) } - /// A Python module implemented in Rust. #[pymodule] fn pyruhvro(m: &Bound<'_, PyModule>) -> PyResult<()> {