Skip to content

Improve Rust reader throughput and resilience #7

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xando merged 1 commit into from
Jun 8, 2026
Merged

Improve Rust reader throughput and resilience #7

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 107 additions & 44 deletions rust/src/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
// Skeleton for the Rust-side BigQuery Storage Read POC.
//
// Wires together gRPC + auth + Arrow IPC decode. Exposed to Python as the
// `_rust.PyReader` class. Iteration yields one `pyarrow.RecordBatch` per
// decoded message from the stream — the Python layer is responsible for any
// rebatching to a target row count.
// `_rust.PyReader` class. Iteration yields groups of decoded
// `pyarrow.RecordBatch` objects drained from the Rust channel.

use std::sync::Arc;
use std::sync::OnceLock;
Expand All @@ -16,8 +15,10 @@ use bytes::Bytes;
use futures::StreamExt;
use pyo3::exceptions::{PyRuntimeError, PyStopIteration};
use pyo3::prelude::*;
use pyo3::types::PyList;
use tokio::runtime::{Builder, Runtime};
use tokio::sync::mpsc;
use tokio::time::{sleep, Duration};
use tonic::metadata::MetadataValue;
use tonic::service::interceptor::InterceptedService;
use tonic::transport::{Channel, ClientTlsConfig, Endpoint};
Expand All @@ -36,6 +37,8 @@ use crate::proto::bqstorage_v1::{

const BQ_STORAGE_ENDPOINT: &str = "https://bigquerystorage.googleapis.com";
const SCOPE: &[&str] = &["https://www.googleapis.com/auth/bigquery.readonly"];
const MAX_DECODING_MESSAGE_SIZE: usize = 256 * 1024 * 1024;
const MAX_READ_ROWS_RETRIES: usize = 3;

fn runtime() -> &'static Runtime {
static RT: OnceLock<Runtime> = OnceLock::new();
Expand Down Expand Up @@ -66,16 +69,18 @@ impl tonic::service::Interceptor for AuthInterceptor {
type AuthedClient =
BigQueryReadClient<InterceptedService<Channel, AuthInterceptor>>;

async fn build_client() -> Result<AuthedClient, String> {
async fn fetch_token() -> Result<Arc<str>, String> {
let provider = gcp_auth::provider()
.await
.map_err(|e| format!("gcp auth: {e}"))?;
let token = provider
.token(SCOPE)
.await
.map_err(|e| format!("token fetch: {e}"))?;
let token_str: Arc<str> = Arc::from(token.as_str());
Ok(Arc::from(token.as_str()))
}

async fn build_client_with_token(token_str: Arc<str>) -> Result<AuthedClient, String> {
let tls = ClientTlsConfig::new().with_native_roots();
let channel = Endpoint::from_static(BQ_STORAGE_ENDPOINT)
.tls_config(tls)
Expand All @@ -87,7 +92,8 @@ async fn build_client() -> Result<AuthedClient, String> {
Ok(BigQueryReadClient::with_interceptor(
channel,
AuthInterceptor { token: token_str },
))
)
.max_decoding_message_size(MAX_DECODING_MESSAGE_SIZE))
}

fn compression_from_str(s: Option<&str>) -> CompressionCodec {
Expand Down Expand Up @@ -144,23 +150,11 @@ async fn create_session(
}

async fn stream_worker(
mut client: AuthedClient,
token: Arc<str>,
stream_name: String,
schema_ipc: Bytes,
tx: mpsc::Sender<Result<RecordBatch, String>>,
) {
let req = ReadRowsRequest {
read_stream: stream_name,
offset: 0,
};
let mut stream = match client.read_rows(req).await {
Ok(s) => s.into_inner(),
Err(e) => {
let _ = tx.send(Err(format!("ReadRows: {e}"))).await;
return;
}
};

// One StreamDecoder per worker: feed the schema bytes once, then each
// gRPC message's record-batch bytes. Avoids the per-message allocate +
// schema-prepend + StreamReader-construction overhead that dominated at
Expand All @@ -175,40 +169,85 @@ async fn stream_worker(
return;
}

while let Some(msg) = stream.next().await {
let response = match msg {
Ok(r) => r,
let mut offset = 0_i64;
let mut retry_count = 0_usize;

'read_attempt: loop {
let mut client = match build_client_with_token(token.clone()).await {
Ok(client) => client,
Err(e) => {
let _ = tx.send(Err(format!("ReadRows stream: {e}"))).await;
let _ = tx.send(Err(format!("client: {e}"))).await;
return;
}
};
let Some(rows) = response.rows else { continue };
use crate::proto::bqstorage_v1::read_rows_response::Rows;
let Rows::ArrowRecordBatch(batch_msg) = rows else { continue };

let batch_slice: &[u8] = batch_msg.serialized_record_batch.as_ref();
let mut buf = ArrowBuffer::from(batch_slice);
while !buf.is_empty() {
match decoder.decode(&mut buf) {
Ok(Some(batch)) => {
if tx.send(Ok(batch)).await.is_err() {
return;
}

let req = ReadRowsRequest {
read_stream: stream_name.clone(),
offset,
};
let mut stream = match client.read_rows(req).await {
Ok(s) => s.into_inner(),
Err(e) => {
if retry_count < MAX_READ_ROWS_RETRIES {
retry_count += 1;
sleep(Duration::from_millis(250 * retry_count as u64)).await;
continue 'read_attempt;
}

let _ = tx.send(Err(format!("ReadRows: {e}"))).await;
return;
}
};

while let Some(msg) = stream.next().await {
let response = match msg {
Ok(r) => {
retry_count = 0;
r
}
Ok(None) => break,
Err(e) => {
let _ = tx.send(Err(format!("arrow ipc decode: {e}"))).await;
if retry_count < MAX_READ_ROWS_RETRIES {
retry_count += 1;
sleep(Duration::from_millis(250 * retry_count as u64)).await;
continue 'read_attempt;
}

let _ = tx.send(Err(format!("ReadRows stream: {e}"))).await;
return;
}
};
let Some(rows) = response.rows else { continue };
use crate::proto::bqstorage_v1::read_rows_response::Rows;
let Rows::ArrowRecordBatch(batch_msg) = rows else { continue };

let batch_slice: &[u8] = batch_msg.serialized_record_batch.as_ref();
let mut buf = ArrowBuffer::from(batch_slice);
while !buf.is_empty() {
match decoder.decode(&mut buf) {
Ok(Some(batch)) => {
offset += batch.num_rows() as i64;

if tx.send(Ok(batch)).await.is_err() {
return;
}
}
Ok(None) => break,
Err(e) => {
let _ = tx.send(Err(format!("arrow ipc decode: {e}"))).await;
return;
}
}
}
}

break;
}
}

#[pyclass(module = "pyarrow.bigquery._rust", unsendable)]
pub struct PyReader {
rx: Option<mpsc::Receiver<Result<RecordBatch, String>>>,
batch_size: usize,
schema_ipc: Vec<u8>,
}

Expand All @@ -228,6 +267,7 @@ impl PyReader {
row_restrictions=None,
max_stream_count=10,
compression=None,
batch_size=100,
channel_capacity=64,
))]
fn new(
Expand All @@ -238,6 +278,7 @@ impl PyReader {
row_restrictions: Option<String>,
max_stream_count: i32,
compression: Option<String>,
batch_size: usize,
channel_capacity: usize,
) -> PyResult<Self> {
let parts: Vec<&str> = source.splitn(3, '.').collect();
Expand All @@ -253,9 +294,10 @@ impl PyReader {
let compression = compression_from_str(compression.as_deref());

let rt = runtime();
let (schema_ipc, streams, client) = py.allow_threads(|| {
let (schema_ipc, streams, token) = py.allow_threads(|| {
rt.block_on(async {
let mut client = build_client().await?;
let token = fetch_token().await?;
let mut client = build_client_with_token(token.clone()).await?;
let session = create_session(
&mut client,
&project,
Expand All @@ -276,7 +318,7 @@ impl PyReader {
};
let streams: Vec<String> =
session.streams.into_iter().map(|s| s.name).collect();
Ok::<_, String>((schema_ipc, streams, client))
Ok::<_, String>((schema_ipc, streams, token))
})
})
.map_err(PyRuntimeError::new_err)?;
Expand All @@ -292,13 +334,14 @@ impl PyReader {
for stream_name in streams {
let tx = tx.clone();
let schema = schema_bytes.clone();
let client = client.clone();
rt.spawn(stream_worker(client, stream_name, schema, tx));
let token = token.clone();
rt.spawn(stream_worker(token, stream_name, schema, tx));
}
drop(tx);

Ok(PyReader {
rx: Some(rx),
batch_size,
schema_ipc,
})
}
Expand All @@ -312,7 +355,7 @@ impl PyReader {
slf
}

/// Pop the next decoded `pyarrow.RecordBatch`, blocking until one is
/// Pop the next decoded `pyarrow.RecordBatch` group, blocking until one is
/// available or all workers have finished.
fn __next__(&mut self, py: Python<'_>) -> PyResult<PyObject> {
let Some(rx) = self.rx.as_mut() else {
Expand All @@ -325,7 +368,27 @@ impl PyReader {
Err(PyStopIteration::new_err(()))
}
Some(Err(e)) => Err(PyRuntimeError::new_err(e)),
Some(Ok(batch)) => batch.to_pyarrow(py),
Some(Ok(batch)) => {
let mut batches = vec![batch];
let mut rows = batches[0].num_rows();
while rows < self.batch_size {
match rx.try_recv() {
Ok(Ok(batch)) => {
rows += batch.num_rows();
batches.push(batch);
}
Ok(Err(e)) => return Err(PyRuntimeError::new_err(e)),
Err(mpsc::error::TryRecvError::Empty) => break,
Err(mpsc::error::TryRecvError::Disconnected) => break,
}
}

let list = PyList::empty_bound(py);
for batch in batches {
list.append(batch.to_pyarrow(py)?)?;
}
Ok(list.into())
}
}
}
}
9 changes: 5 additions & 4 deletions src/pyarrow/bigquery/read.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,8 +345,9 @@ def _enter_rust(self):
self.project,
columns=self.columns,
row_restrictions=self.row_restrictions,
max_stream_count=self.worker_count * 3,
max_stream_count=self.worker_count,
compression=self.compression,
batch_size=self.batch_size,
)
self.schema = pa.ipc.read_schema(pa.py_buffer(self._rust_reader.schema_ipc()))
self._rust_pending: list[pa.RecordBatch] = []
Expand Down Expand Up @@ -418,12 +419,12 @@ def _raise_if_workers_failed(self):
def _next_rust(self) -> pa.Table:
while not self._rust_exhausted and self._rust_pending_rows < self.batch_size:
try:
batch = next(self._rust_reader)
batches = next(self._rust_reader)
except StopIteration:
self._rust_exhausted = True
break
self._rust_pending.append(batch)
self._rust_pending_rows += batch.num_rows
self._rust_pending.extend(batches)
self._rust_pending_rows += sum(batch.num_rows for batch in batches)

if not self._rust_pending:
raise StopIteration
Expand Down
Loading