arrow-bigquery 0.3.0__tar.gz → 0.3.2__tar.gz

{arrow_bigquery-0.3.0 → arrow_bigquery-0.3.2}/PKG-INFO

@@ -1,10 +1,10 @@
 Metadata-Version: 2.4
 Name: arrow-bigquery
-Version: 0.3.0
+Version: 0.3.2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Dist: google-auth>=2.40.3
 Requires-Dist: requests>=2.33.0,<3.0.0
 Requires-Dist: cachetools>=5.0.0
-Requires-Python: >=3.10
+Requires-Python: >=3.11

{arrow_bigquery-0.3.0 → arrow_bigquery-0.3.2}/arrow_bigquery/_read_bigquery.py

@@ -80,7 +80,7 @@ class Client:
         *,
         arrow_buffer_compression: str = "lz4frame",
         dry_run: bool = False,
-        maintain_order: bool = False,
+        maintain_order: bool = True,
         use_query_cache: bool | None = None,
     ) -> arrow_bigquery._native.ArrowStreamExporter:
         return self._client.read_query(

arrow_bigquery-0.3.2/arrow_bigquery/core/version.py

@@ -0,0 +1 @@
+__version__ = "0.3.2"

{arrow_bigquery-0.3.0 → arrow_bigquery-0.3.2}/crates/arrow-bigquery/src/bigquery_read_stream.rs

@@ -198,7 +198,6 @@ pub(crate) struct ReadStreamOptions {
     pub(crate) first_page_check: Option<FirstPageCheck>,
     pub(crate) tx: tokio::sync::mpsc::Sender<Result<RecordBatch, BigQueryError>>,
     pub(crate) retry_config: StreamRetryConfig,
-    pub(crate) decode_pool: Arc<rayon::ThreadPool>,
     pub(crate) maintain_order: bool,
     pub(crate) cancel_tx: Arc<tokio::sync::watch::Sender<bool>>,
     pub(crate) cancel_rx: tokio::sync::watch::Receiver<bool>,
@@ -211,7 +210,6 @@ pub(crate) struct ReadQueryStreamOptions {
     pub(crate) metadata: Arc<StreamMetadata>,
     pub(crate) tx: tokio::sync::mpsc::Sender<Result<RecordBatch, BigQueryError>>,
     pub(crate) retry_config: StreamRetryConfig,
-    pub(crate) decode_pool: Arc<rayon::ThreadPool>,
     pub(crate) maintain_order: bool,
     pub(crate) cancel_tx: Arc<tokio::sync::watch::Sender<bool>>,
     pub(crate) cancel_rx: tokio::sync::watch::Receiver<bool>,
@@ -241,7 +239,6 @@ impl ReadStream {
             first_page_check,
             tx,
             retry_config,
-            decode_pool: _,
             maintain_order: _,
             cancel_tx,
             cancel_rx,
@@ -576,7 +573,7 @@ impl ReadStream {
 
 /// Orchestrates stream connection, consumption, and mid-stream reconnections via [`ReadStream::step`].
 pub(crate) async fn read_stream(options: ReadStreamOptions) -> TerminationReason {
-    let decode_queue = DecodeQueue::new(Arc::clone(&options.decode_pool), options.maintain_order);
+    let decode_queue = DecodeQueue::new(options.maintain_order);
     let mut reader = ReadStream::new(options, decode_queue);
     reader.run().await
 }
@@ -592,12 +589,11 @@ pub(crate) async fn read_query_stream(options: ReadQueryStreamOptions) {
         metadata,
         tx,
         retry_config,
-        decode_pool,
         maintain_order,
         cancel_tx,
         mut cancel_rx,
     } = options;
-    let decode_queue = DecodeQueue::new(Arc::clone(&decode_pool), maintain_order);
+    let decode_queue = DecodeQueue::new(maintain_order);
 
     if !enqueue_decode_batch(
         &decode_queue,
@@ -623,7 +619,6 @@ pub(crate) async fn read_query_stream(options: ReadQueryStreamOptions) {
                 first_page_check: None,
                 tx,
                 retry_config,
-                decode_pool,
                 maintain_order,
                 cancel_tx,
                 cancel_rx,
@@ -651,7 +646,6 @@ pub(crate) mod tests {
 
     use super::*;
     use crate::bigquery_read_retry::RetryableErrors;
-    use crate::decode::default_decode_pool;
 
     #[derive(Debug, Default)]
     pub(crate) struct RecordedReadRows {
@@ -906,7 +900,6 @@ pub(crate) mod tests {
             first_page_check: None,
             tx,
             retry_config: test_retry_config(3),
-            decode_pool: default_decode_pool(),
             maintain_order: true,
             cancel_tx: Arc::new(cancel_tx),
             cancel_rx,
@@ -956,7 +949,6 @@ pub(crate) mod tests {
             first_page_check: None,
             tx,
             retry_config: test_retry_config(3),
-            decode_pool: default_decode_pool(),
             maintain_order: false,
             cancel_tx: Arc::new(cancel_tx),
             cancel_rx,
@@ -999,7 +991,6 @@ pub(crate) mod tests {
             first_page_check: None,
             tx,
             retry_config: test_retry_config(3),
-            decode_pool: default_decode_pool(),
             maintain_order: false,
             cancel_tx: Arc::clone(&cancel_tx),
             cancel_rx,
@@ -1054,7 +1045,6 @@ pub(crate) mod tests {
             first_page_check: None,
             tx,
             retry_config: test_retry_config(3),
-            decode_pool: default_decode_pool(),
             maintain_order: false,
             cancel_tx: Arc::clone(&cancel_tx),
             cancel_rx,
@@ -1112,7 +1102,6 @@ pub(crate) mod tests {
             first_page_check: None,
             tx,
             retry_config: test_retry_config(3),
-            decode_pool: default_decode_pool(),
             maintain_order: false,
             cancel_tx: Arc::clone(&cancel_tx),
             cancel_rx,
@@ -1143,15 +1132,17 @@ pub(crate) mod tests {
         )
     }
 
+    type TestReadStream = (
+        ReadStream,
+        Arc<Mutex<Vec<ReadRowsRequest>>>,
+        tokio::sync::mpsc::Receiver<Result<RecordBatch, BigQueryError>>,
+    );
+
     fn make_test_read_stream(
         streams: Vec<MockStreamResult>,
         metadata: Option<Arc<StreamMetadata>>,
         retry_config: StreamRetryConfig,
-    ) -> (
-        ReadStream,
-        Arc<Mutex<Vec<ReadRowsRequest>>>,
-        tokio::sync::mpsc::Receiver<Result<RecordBatch, BigQueryError>>,
-    ) {
+    ) -> TestReadStream {
         let requests = Arc::new(Mutex::new(vec![]));
         let mock_client = MockClient {
             requests: Arc::clone(&requests),
@@ -1161,8 +1152,7 @@ pub(crate) mod tests {
         let read_rows = client.read_rows().set_read_stream("streams/1");
         let (tx, rx) = tokio::sync::mpsc::channel(10);
         let (cancel_tx, cancel_rx) = tokio::sync::watch::channel(false);
-        let decode_pool = default_decode_pool();
-        let decode_queue = DecodeQueue::new(Arc::clone(&decode_pool), false);
+        let decode_queue = DecodeQueue::new(false);
         let reader = ReadStream::new(
             ReadStreamOptions {
                 read_rows,
@@ -1172,7 +1162,6 @@ pub(crate) mod tests {
                 first_page_check: None,
                 tx,
                 retry_config,
-                decode_pool,
                 maintain_order: false,
                 cancel_tx: Arc::new(cancel_tx),
                 cancel_rx,

{arrow_bigquery-0.3.0 → arrow_bigquery-0.3.2}/crates/arrow-bigquery/src/decode.rs

@@ -1,6 +1,6 @@
 use std::collections::VecDeque;
 use std::io::Cursor;
-use std::sync::{Arc, LazyLock, Mutex};
+use std::sync::{Arc, Mutex};
 
 use polars_arrow::io::ipc::read::{
     read_stream_metadata, StreamMetadata, StreamReader, StreamState,
@@ -9,24 +9,6 @@ use polars_arrow::record_batch::RecordBatch;
 
 use crate::BigQueryError;
 
-static DEFAULT_DECODE_POOL: LazyLock<Arc<rayon::ThreadPool>> = LazyLock::new(|| {
-    let num_threads = std::thread::available_parallelism()
-        .map(|n| n.get())
-        .unwrap_or(1);
-    let pool = rayon::ThreadPoolBuilder::new()
-        .num_threads(num_threads)
-        .thread_name(|i| format!("arrow-bq-decode-{i}"))
-        .build()
-        .expect("failed to build Arrow IPC decode thread pool");
-    Arc::new(pool)
-});
-
-/// Returns a handle to the shared bounded thread pool used for CPU-bound
-/// Arrow IPC (LZ4/ZSTD) decompression.
-pub(crate) fn default_decode_pool() -> Arc<rayon::ThreadPool> {
-    Arc::clone(&DEFAULT_DECODE_POOL)
-}
-
 /// Parses IPC-serialized Arrow schema bytes into [`StreamMetadata`].
 pub(crate) fn parse_arrow_schema(schema_bytes: &[u8]) -> Result<StreamMetadata, BigQueryError> {
     let mut cursor = Cursor::new(schema_bytes);
@@ -150,14 +132,12 @@ struct OrderedQueueState {
 /// A FIFO decode queue that keeps all decode jobs for a `maintain_order = true`
 /// read in a single worker queue without work-stealing across workers.
 pub(crate) struct OrderedDecodeQueue {
-    decode_pool: Arc<rayon::ThreadPool>,
     state: Mutex<OrderedQueueState>,
 }
 
 impl OrderedDecodeQueue {
-    pub(crate) fn new(decode_pool: Arc<rayon::ThreadPool>) -> Arc<Self> {
+    pub(crate) fn new() -> Arc<Self> {
         Arc::new(Self {
-            decode_pool,
             state: Mutex::new(OrderedQueueState {
                 queue: VecDeque::new(),
                 running: false,
@@ -179,7 +159,7 @@ impl OrderedDecodeQueue {
 
         if should_spawn {
             let queue = Arc::clone(self);
-            self.decode_pool.spawn_fifo(move || {
+            rayon::spawn_fifo(move || {
                 queue.drain();
             });
         }
@@ -204,25 +184,25 @@ impl OrderedDecodeQueue {
 
 #[derive(Clone)]
 pub(crate) enum DecodeQueue {
-    /// Uses Rayon's work-stealing queue when `maintain_order = false`.
-    WorkStealing(Arc<rayon::ThreadPool>),
+    /// Uses Rayon's global work-stealing pool when `maintain_order = false`.
+    WorkStealing,
     /// Keeps all decode jobs in a single sequential queue when `maintain_order = true`.
     Ordered(Arc<OrderedDecodeQueue>),
 }
 
 impl DecodeQueue {
-    pub(crate) fn new(decode_pool: Arc<rayon::ThreadPool>, maintain_order: bool) -> Self {
+    pub(crate) fn new(maintain_order: bool) -> Self {
         if maintain_order {
-            Self::Ordered(OrderedDecodeQueue::new(decode_pool))
+            Self::Ordered(OrderedDecodeQueue::new())
         } else {
-            Self::WorkStealing(decode_pool)
+            Self::WorkStealing
         }
     }
 
     pub(crate) fn push(&self, job: DecodeJob) {
         match self {
-            Self::WorkStealing(pool) => {
-                pool.spawn_fifo(move || job.run());
+            Self::WorkStealing => {
+                rayon::spawn_fifo(move || job.run());
             },
             Self::Ordered(queue) => {
                 queue.push(job);
@@ -240,13 +220,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_ordered_decode_queue_preserves_batch_order() {
-        let pool = Arc::new(
-            rayon::ThreadPoolBuilder::new()
-                .num_threads(4)
-                .build()
-                .unwrap(),
-        );
-        let queue = DecodeQueue::new(pool, true);
+        let queue = DecodeQueue::new(true);
 
         // Earlier batches are much larger than later ones so work-stealing would finish
         // the smaller batches first if they were not kept in the worker's ordered queue.
@@ -302,13 +276,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_work_stealing_decode_queue_decodes_all_batches() {
-        let pool = Arc::new(
-            rayon::ThreadPoolBuilder::new()
-                .num_threads(4)
-                .build()
-                .unwrap(),
-        );
-        let queue = DecodeQueue::new(pool, false);
+        let queue = DecodeQueue::new(false);
 
         let sizes = [20_000, 10, 15_000, 10, 10_000, 10];
         let mut schema_bytes = Vec::new();
@@ -358,13 +326,7 @@ mod tests {
 
     #[tokio::test]
     async fn test_decode_job_error_signals_cancellation_and_pushes_err() {
-        let pool = Arc::new(
-            rayon::ThreadPoolBuilder::new()
-                .num_threads(2)
-                .build()
-                .unwrap(),
-        );
-        let queue = DecodeQueue::new(pool, true);
+        let queue = DecodeQueue::new(true);
 
         let (schema_bytes, valid_batch) = create_test_arrow_payload_with_range(0, 5);
         let metadata = Arc::new(parse_arrow_schema(&schema_bytes).unwrap());