quickhouse 0.2.0__tar.gz → 0.2.3__tar.gz

{quickhouse-0.2.0 → quickhouse-0.2.3}/Cargo.lock

@@ -644,6 +644,8 @@ dependencies = [
  "compression-core",
  "flate2",
  "memchr",
+ "zstd",
+ "zstd-safe",
 ]
 
 [[package]]
@@ -2504,15 +2506,15 @@ dependencies = [
 
 [[package]]
 name = "quickhouse-core"
-version = "0.2.0"
+version = "0.2.3"
 dependencies = [
  "anyhow",
  "arrow",
  "arrow-array",
  "arrow-schema",
+ "async-compression",
  "bytes",
  "chrono",
- "flate2",
  "futures",
  "google-cloud-bigquery",
  "google-cloud-googleapis",
@@ -2527,13 +2529,14 @@ dependencies = [
  "tokio",
  "tokio-postgres",
  "tokio-postgres-rustls",
+ "tokio-util",
  "tracing",
  "webpki-roots 0.26.11",
 ]
 
 [[package]]
 name = "quickhouse-py"
-version = "0.2.0"
+version = "0.2.3"
 dependencies = [
  "pyo3",
  "quickhouse-core",
@@ -3368,9 +3371,9 @@ dependencies = [
 
 [[package]]
 name = "tokio"
-version = "1.52.4"
+version = "1.53.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "317fafbbe3f02fc663dad00ea6186197de963cd4190e86a26d8d0fae095539af"
+checksum = "d988bcd52dbe076d3d46903332f58c912b87a2c49b1428419a5845154762ffee"
 dependencies = [
  "bytes",
  "libc",
@@ -3384,9 +3387,9 @@ dependencies = [
 
 [[package]]
 name = "tokio-macros"
-version = "2.7.0"
+version = "2.7.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496"
+checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba"
 dependencies = [
  "proc-macro2",
  "quote",

{quickhouse-0.2.0 → quickhouse-0.2.3}/Cargo.toml

@@ -3,7 +3,7 @@ resolver = "2"
 members = ["crates/quickhouse-core", "crates/quickhouse-py"]
 
 [workspace.package]
-version = "0.2.0"
+version = "0.2.3"
 edition = "2021"
 license = "MIT"
 authors = ["quickhouse contributors"]
@@ -19,6 +19,8 @@ arrow-array = "53"
 arrow-schema = "53"
 reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "stream", "gzip"] }
 bytes = "1"
+tokio-util = { version = "0.7", features = ["io"] }
+async-compression = { version = "0.4", features = ["tokio", "gzip", "zstd"] }
 thiserror = "1"
 anyhow = "1"
 tracing = "0.1"
@@ -27,7 +29,6 @@ serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 chrono = { version = "0.4", default-features = false, features = ["std"] }
 uuid = { version = "1", features = ["v4"] }
-flate2 = "1"
 rustls = { version = "0.23", default-features = false, features = ["ring", "logging", "std", "tls12"] }
 tokio-postgres-rustls = "0.13"
 webpki-roots = "0.26"

{quickhouse-0.2.0 → quickhouse-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: quickhouse
-Version: 0.2.0
+Version: 0.2.3
 Classifier: Development Status :: 3 - Alpha
 Classifier: Intended Audience :: Developers
 Classifier: License :: OSI Approved :: MIT License
@@ -73,9 +73,11 @@ either way.
 | --- | --- |
 | Deserialization | PostgreSQL: binary `COPY` decoded straight into Arrow in Rust. MySQL: `mysql_async`'s typed binary-protocol rows appended straight into Arrow builders. BigQuery: the Storage Read API's wire format is already Arrow. Either way, no per-row Python objects. |
 | Parallelism | Postgres/MySQL: table split into key ranges, one connection + Tokio task per partition. BigQuery: `parallelism` is passed as a stream-count hint to BigQuery's own server-side parallel preparation (see the BigQuery note below). |
-| Memory | Batches flushed every `batch_rows` rows *or* `batch_bytes` (whichever hits first) and streamed to ClickHouse — RSS stays flat regardless of table size, and wide-row tables can't blow past the byte budget the way a rows-only cap would |
+| Pipelining | Decoding overlaps with uploading: each finished batch's insert is spawned as a background task, so a partition keeps reading/decoding while previous batches are still being compressed and sent, instead of stalling on each HTTP round-trip. |
+| Memory | A single hard ceiling, `max_memory_bytes`, bounds **total** in-flight batch memory (across every partition and every upload in flight), measured against each batch's real Arrow allocation. When the ceiling is reached, decoding blocks (backpressure) until uploads drain — so peak RSS stays bounded regardless of `parallelism`, row width, or partition skew. `batch_rows`/`batch_bytes` separately control how big each individual batch is. |
 | GIL | Entire transfer runs inside `Python::allow_threads`; the GIL is only touched for `on_progress` |
-| Insert | Arrow IPC stream ingested by ClickHouse's native `ArrowStream` format, gzip on the wire |
+| Insert | Arrow IPC stream ingested by ClickHouse's native `ArrowStream` format, streamed to the wire with zstd (default), gzip, or no compression — the compressed body is produced incrementally rather than buffered in full. |
+| Resilience | Each insert retries transient failures (dropped/reset connections, timeouts, HTTP 5xx/429) with exponential backoff — up to 4 attempts — so a single network blip over a long WAN transfer doesn't abort the whole run. Deterministic 4xx errors (bad SQL, auth) fail fast without retrying. Retry is at-least-once: a lost ack after a committed batch can duplicate that one batch — harmless in incremental mode (`ReplacingMergeTree` dedupes by key), possible dupes for one batch in full-refresh into a plain `MergeTree`. |
 
 ## Sync modes
 
@@ -170,8 +172,9 @@ cargo test -p quickhouse-core
 | `create_if_missing` | Auto-run generated `CREATE TABLE` when the destination is absent |
 | `engine`, `order_by`, `partition_by`, `primary_key` | DDL knobs (sensible defaults per mode) |
 | `parallelism` | Number of concurrent partition streams |
-| `batch_rows` | Max rows per Arrow batch / insert flush (memory vs round-trips) |
-| `batch_bytes` | Max estimated source bytes per batch before flushing, even under `batch_rows` — default 4 MiB; caps peak memory for wide-row tables. `0` disables (row count alone decides) |
+| `batch_rows` | Max rows in each Arrow batch / insert — a per-batch granularity knob (round-trips vs. per-insert overhead), **not** the memory ceiling |
+| `batch_bytes` | Also cap each individual batch at this many estimated source bytes, even under `batch_rows` — default 4 MiB; keeps a single wide-row batch from growing large. `0` disables (row count alone sizes the batch) |
+| `max_memory_bytes` | **The memory ceiling.** Hard cap on total in-flight Arrow batch memory across all partitions and all in-flight uploads, measured against each batch's real allocation. Decoding blocks (backpressure) when reached, so peak RSS stays bounded independent of `parallelism`/row width. Default 512 MiB; `0` = unbounded |
 | `partition_column` | Integer column to range-split on (defaults to first `key`) |
 | `type_overrides` | Per-column ClickHouse type (e.g. `{"qty": "Decimal(18, 3)"}`) |
 | `rename` | Source → destination column renames |

{quickhouse-0.2.0 → quickhouse-0.2.3}/README.md

@@ -43,9 +43,11 @@ either way.
 | --- | --- |
 | Deserialization | PostgreSQL: binary `COPY` decoded straight into Arrow in Rust. MySQL: `mysql_async`'s typed binary-protocol rows appended straight into Arrow builders. BigQuery: the Storage Read API's wire format is already Arrow. Either way, no per-row Python objects. |
 | Parallelism | Postgres/MySQL: table split into key ranges, one connection + Tokio task per partition. BigQuery: `parallelism` is passed as a stream-count hint to BigQuery's own server-side parallel preparation (see the BigQuery note below). |
-| Memory | Batches flushed every `batch_rows` rows *or* `batch_bytes` (whichever hits first) and streamed to ClickHouse — RSS stays flat regardless of table size, and wide-row tables can't blow past the byte budget the way a rows-only cap would |
+| Pipelining | Decoding overlaps with uploading: each finished batch's insert is spawned as a background task, so a partition keeps reading/decoding while previous batches are still being compressed and sent, instead of stalling on each HTTP round-trip. |
+| Memory | A single hard ceiling, `max_memory_bytes`, bounds **total** in-flight batch memory (across every partition and every upload in flight), measured against each batch's real Arrow allocation. When the ceiling is reached, decoding blocks (backpressure) until uploads drain — so peak RSS stays bounded regardless of `parallelism`, row width, or partition skew. `batch_rows`/`batch_bytes` separately control how big each individual batch is. |
 | GIL | Entire transfer runs inside `Python::allow_threads`; the GIL is only touched for `on_progress` |
-| Insert | Arrow IPC stream ingested by ClickHouse's native `ArrowStream` format, gzip on the wire |
+| Insert | Arrow IPC stream ingested by ClickHouse's native `ArrowStream` format, streamed to the wire with zstd (default), gzip, or no compression — the compressed body is produced incrementally rather than buffered in full. |
+| Resilience | Each insert retries transient failures (dropped/reset connections, timeouts, HTTP 5xx/429) with exponential backoff — up to 4 attempts — so a single network blip over a long WAN transfer doesn't abort the whole run. Deterministic 4xx errors (bad SQL, auth) fail fast without retrying. Retry is at-least-once: a lost ack after a committed batch can duplicate that one batch — harmless in incremental mode (`ReplacingMergeTree` dedupes by key), possible dupes for one batch in full-refresh into a plain `MergeTree`. |
 
 ## Sync modes
 
@@ -140,8 +142,9 @@ cargo test -p quickhouse-core
 | `create_if_missing` | Auto-run generated `CREATE TABLE` when the destination is absent |
 | `engine`, `order_by`, `partition_by`, `primary_key` | DDL knobs (sensible defaults per mode) |
 | `parallelism` | Number of concurrent partition streams |
-| `batch_rows` | Max rows per Arrow batch / insert flush (memory vs round-trips) |
-| `batch_bytes` | Max estimated source bytes per batch before flushing, even under `batch_rows` — default 4 MiB; caps peak memory for wide-row tables. `0` disables (row count alone decides) |
+| `batch_rows` | Max rows in each Arrow batch / insert — a per-batch granularity knob (round-trips vs. per-insert overhead), **not** the memory ceiling |
+| `batch_bytes` | Also cap each individual batch at this many estimated source bytes, even under `batch_rows` — default 4 MiB; keeps a single wide-row batch from growing large. `0` disables (row count alone sizes the batch) |
+| `max_memory_bytes` | **The memory ceiling.** Hard cap on total in-flight Arrow batch memory across all partitions and all in-flight uploads, measured against each batch's real allocation. Decoding blocks (backpressure) when reached, so peak RSS stays bounded independent of `parallelism`/row width. Default 512 MiB; `0` = unbounded |
 | `partition_column` | Integer column to range-split on (defaults to first `key`) |
 | `type_overrides` | Per-column ClickHouse type (e.g. `{"qty": "Decimal(18, 3)"}`) |
 | `rename` | Source → destination column renames |

{quickhouse-0.2.0 → quickhouse-0.2.3}/crates/quickhouse-core/Cargo.toml

@@ -17,12 +17,13 @@ arrow-array.workspace = true
 arrow-schema.workspace = true
 reqwest.workspace = true
 bytes.workspace = true
+tokio-util.workspace = true
+async-compression.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 chrono.workspace = true
-flate2.workspace = true
 rustls.workspace = true
 tokio-postgres-rustls.workspace = true
 webpki-roots.workspace = true

{quickhouse-0.2.0 → quickhouse-0.2.3}/crates/quickhouse-core/src/config.rs

@@ -73,7 +73,7 @@ pub struct ClickHouseConfig {
     pub database: String,
     pub user: String,
     pub password: String,
-    /// `"none" | "gzip"` — HTTP body compression for inserts.
+    /// `"none" | "gzip" | "zstd"` — HTTP body compression for inserts.
     pub compression: Compression,
 }
 
@@ -81,6 +81,7 @@ pub struct ClickHouseConfig {
 pub enum Compression {
     None,
     Gzip,
+    Zstd,
 }
 
 /// Full-refresh reloads everything; Incremental appends rows past a watermark.
@@ -118,13 +119,26 @@ pub struct TransferConfig {
 
     // ---- parallelism / batching ----
     pub parallelism: usize,
+    /// Per-batch granularity: flush a RecordBatch once it reaches this many
+    /// rows. Controls how big each individual insert is (a throughput/overhead
+    /// knob), NOT the overall memory ceiling — that's `max_memory_bytes`.
     pub batch_rows: usize,
-    /// Flush a RecordBatch once its accumulated (estimated) source bytes
-    /// reach this many, even if `batch_rows` hasn't been hit yet — caps peak
-    /// memory for wide-row tables, where row-count alone is a poor proxy for
-    /// actual bytes buffered. `0` disables the byte-based trigger (row count
-    /// alone decides, as before this option existed).
+    /// Per-batch granularity: also flush once a batch's accumulated (estimated)
+    /// source bytes reach this many, even if `batch_rows` hasn't been hit yet,
+    /// so a single batch of wide rows doesn't grow unbounded. `0` disables this
+    /// per-batch byte cap (row count alone decides batch size). This bounds one
+    /// *batch*; the total in-flight memory across all partitions and in-flight
+    /// inserts is bounded separately by `max_memory_bytes`.
     pub batch_bytes: usize,
+    /// Hard ceiling on total in-flight Arrow batch memory across the whole
+    /// transfer — every partition's decoded-but-not-yet-sent batches plus all
+    /// batches currently being uploaded. Enforced against each batch's real
+    /// `RecordBatch::get_array_memory_size()`, so it holds regardless of
+    /// `parallelism`, row width, or partition skew. When the ceiling is
+    /// reached, decoding blocks (backpressure) until in-flight inserts drain.
+    /// `0` disables the ceiling (unbounded — memory then scales with
+    /// parallelism and batch size, the pre-`max_memory_bytes` behavior).
+    pub max_memory_bytes: usize,
     /// Column used to split the table into parallel range partitions.
     /// Defaults to the first `key` column, else the sync falls back to a single stream.
     pub partition_column: Option<String>,
@@ -169,6 +183,14 @@ impl TransferConfig {
         if self.batch_rows == 0 {
             return Err(EtlError::config("batch_rows must be >= 1"));
         }
+        // A non-zero ceiling must at least admit a single batch's worth of
+        // rows-ish of memory; guard against pathologically tiny values that
+        // would stall every transfer. (0 = unbounded, always allowed.)
+        if self.max_memory_bytes != 0 && self.max_memory_bytes < 64 * 1024 {
+            return Err(EtlError::config(
+                "max_memory_bytes must be 0 (unbounded) or >= 65536",
+            ));
+        }
         Ok(())
     }
 }

{quickhouse-0.2.0 → quickhouse-0.2.3}/crates/quickhouse-core/src/ddl.rs

@@ -138,6 +138,7 @@ mod tests {
             parallelism: 4,
             batch_rows: 1000,
             batch_bytes: 0,
+            max_memory_bytes: 0,
             partition_column: None,
             type_overrides: HashMap::new(),
             rename: HashMap::new(),

{quickhouse-0.2.0 → quickhouse-0.2.3}/crates/quickhouse-core/src/lib.rs

@@ -15,6 +15,7 @@ pub mod decode;
 pub mod decode_bigquery;
 pub mod decode_mysql;
 pub mod error;
+pub mod memory;
 pub mod sink;
 pub mod source;
 pub mod sync;

quickhouse-0.2.3/crates/quickhouse-core/src/memory.rs

@@ -0,0 +1,146 @@
+//! Process-wide memory budget for in-flight Arrow batches.
+//!
+//! The transfer pipeline decodes batches on one side and uploads them to
+//! ClickHouse concurrently on the other. Without a ceiling, a fast source and
+//! a slow sink would let decoded-but-not-yet-sent batches pile up unbounded.
+//!
+//! [`MemoryBudget`] is a byte-denominated semaphore (1 permit = 1 byte): a
+//! batch reserves permits equal to its real [`RecordBatch::get_array_memory_size`]
+//! before it may be handed to a send task, and releases them when the send
+//! completes. This makes peak in-flight memory a single honest number
+//! (`max_memory_bytes`) that holds regardless of parallelism, row width, or
+//! partition skew — unlike a per-batch row/byte heuristic, whose true peak is
+//! an emergent `parallelism × batch_size`.
+//!
+//! [`RecordBatch::get_array_memory_size`]: arrow_array::RecordBatch::get_array_memory_size
+
+use std::sync::Arc;
+
+use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+
+/// tokio's `Semaphore` caps total permits at `usize::MAX >> 3`, and
+/// `acquire_many_owned` takes a `u32`. Batches are MiB-scale so a single
+/// reservation never approaches this, but clamp defensively so an absurd
+/// size can never panic or wrap.
+const MAX_PERMITS: u32 = u32::MAX;
+
+/// A shared ceiling on total in-flight Arrow batch memory.
+///
+/// Cloneable and cheap to share across tasks (an `Arc` inside). `None`-like
+/// "unbounded" is represented by constructing with `max_bytes == 0`.
+#[derive(Clone)]
+pub struct MemoryBudget {
+    /// `None` = unbounded (no accounting at all).
+    sem: Option<Arc<Semaphore>>,
+    /// The configured ceiling, retained so a single oversized batch can be
+    /// clamped to "reserve the whole budget" rather than deadlock.
+    total: usize,
+}
+
+impl MemoryBudget {
+    /// Create a budget of `max_bytes`. `0` means unbounded.
+    pub fn new(max_bytes: usize) -> Self {
+        let sem = if max_bytes == 0 {
+            None
+        } else {
+            // Semaphore permit count is a usize; cap at the tokio maximum.
+            let permits = max_bytes.min(Semaphore::MAX_PERMITS);
+            Some(Arc::new(Semaphore::new(permits)))
+        };
+        MemoryBudget {
+            sem,
+            total: max_bytes,
+        }
+    }
+
+    /// Whether this budget actually enforces a ceiling.
+    pub fn is_bounded(&self) -> bool {
+        self.sem.is_some()
+    }
+
+    /// Reserve room for a batch of `size` bytes, awaiting until enough is free.
+    ///
+    /// Returns a permit whose `Drop` releases the reservation — hold it for as
+    /// long as the batch occupies memory (i.e. until its upload finishes).
+    /// A batch larger than the whole budget is clamped to the full budget so
+    /// it still runs (alone) rather than deadlocking forever.
+    pub async fn reserve(&self, size: usize) -> Reservation {
+        match &self.sem {
+            None => Reservation { _permit: None },
+            Some(sem) => {
+                let want = size.min(self.total).min(MAX_PERMITS as usize).max(1) as u32;
+                // acquire_many_owned only errors if the semaphore is closed,
+                // which we never do; unwrap is safe for the lifetime of self.
+                let permit = sem
+                    .clone()
+                    .acquire_many_owned(want)
+                    .await
+                    .expect("memory budget semaphore closed unexpectedly");
+                Reservation {
+                    _permit: Some(permit),
+                }
+            }
+        }
+    }
+}
+
+/// An outstanding memory reservation. Dropping it returns the bytes to the
+/// budget. Opaque on purpose — callers only need to keep it alive.
+pub struct Reservation {
+    _permit: Option<OwnedSemaphorePermit>,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[tokio::test]
+    async fn unbounded_never_blocks() {
+        let budget = MemoryBudget::new(0);
+        assert!(!budget.is_bounded());
+        // Many large reservations resolve immediately and independently.
+        let _a = budget.reserve(10_000_000_000).await;
+        let _b = budget.reserve(10_000_000_000).await;
+    }
+
+    #[tokio::test]
+    async fn releases_on_drop() {
+        let budget = MemoryBudget::new(1000);
+        {
+            let _r = budget.reserve(1000).await; // takes the whole budget
+                                                 // A second full reservation would block here; instead we drop first.
+        }
+        // After drop, the full budget is available again — this must not hang.
+        let _r2 = budget.reserve(1000).await;
+    }
+
+    #[tokio::test]
+    async fn oversized_batch_still_proceeds() {
+        // A batch bigger than the entire budget must not deadlock; it is
+        // clamped to the whole budget and runs alone.
+        let budget = MemoryBudget::new(1000);
+        let _r = budget.reserve(5000).await; // clamped to 1000, resolves
+    }
+
+    #[tokio::test]
+    async fn blocks_until_released() {
+        use std::time::Duration;
+        use tokio::time::timeout;
+
+        let budget = MemoryBudget::new(1000);
+        let first = budget.reserve(1000).await; // exhausts the budget
+
+        // A second reservation cannot complete while `first` is held.
+        let second = budget.reserve(1000);
+        assert!(
+            timeout(Duration::from_millis(50), second).await.is_err(),
+            "reservation should block while the budget is fully held"
+        );
+
+        // Releasing the first frees the budget for a waiter.
+        drop(first);
+        let _third = timeout(Duration::from_millis(500), budget.reserve(1000))
+            .await
+            .expect("reservation should proceed after the budget is released");
+    }
+}