feat: Query::fetch_bytes (#182)

Co-authored-by: Paul Loyd <pavelko95@gmail.com>

Author: slvrtrn

.github/workflows/ci.yml

@@ -32,6 +32,9 @@ jobs:
     - run: rustup toolchain install ${{ env.MSRV }} --profile minimal
     - run: rustup override set ${{ env.MSRV }}
     - run: rustup show active-toolchain -v
+    - run: cargo update -p native-tls --precise 0.2.13 # 0.2.14 requires rustc 1.80
+    - run: cargo update -p litemap --precise 0.7.4 # 0.7.5 requires rustc 1.81
+    - run: cargo update -p zerofrom --precise 0.1.5 # 0.1.6 requires rustc 1.81
     - run: cargo build
     - run: cargo build --no-default-features
     - run: cargo build --features uuid,time,chrono

Cargo.toml

@@ -10,6 +10,7 @@ license = "MIT OR Apache-2.0"
 readme = "README.md"
 edition = "2021"
 # update `derive/Cargo.toml` and CI if changed
+# TODO: after bumping to v1.80, remove `--precise` in the "msrv" CI job
 rust-version = "1.73.0"
 
 [lints.rust]
@@ -70,6 +71,8 @@ uuid = ["dep:uuid"]
 time = ["dep:time"]
 lz4 = ["dep:lz4_flex", "dep:cityhash-rs"]
 chrono = ["dep:chrono"]
+futures03 = []
+
 ##  TLS
 native-tls = ["dep:hyper-tls"]
 # ext: native-tls-alpn

benches/README.md

@@ -23,8 +23,7 @@ Then upload the `perf.script` file to [Firefox Profiler](https://profiler.firefo
 
 These benchmarks are run against a real ClickHouse server, so it must be started:
 ```bash
-docker run -d -p 8123:8123 -p 9000:9000 --name ch clickhouse/clickhouse-server
-
+docker compose up -d
 cargo bench --bench <case>
 ```
 

benches/select.rs

@@ -14,7 +14,10 @@ use hyper::{
 };
 use serde::Deserialize;
 
-use clickhouse::{error::Result, Client, Compression, Row};
+use clickhouse::{
+    error::{Error, Result},
+    Client, Compression, Row,
+};
 
 mod common;
 
@@ -52,36 +55,80 @@ fn select(c: &mut Criterion) {
     let _server = common::start_server(addr, move |req| serve(req, chunk.clone()));
     let runner = common::start_runner();
 
-    #[allow(dead_code)]
-    #[derive(Debug, Row, Deserialize)]
+    #[derive(Default, Debug, Row, Deserialize)]
     struct SomeRow {
         a: u64,
         b: i64,
         c: i32,
         d: u32,
     }
 
-    async fn run(client: Client, iters: u64) -> Result<Duration> {
+    async fn select_rows(client: Client, iters: u64) -> Result<Duration> {
+        let mut sum = SomeRow::default();
         let start = Instant::now();
         let mut cursor = client
             .query("SELECT ?fields FROM some")
             .fetch::<SomeRow>()?;
 
         for _ in 0..iters {
-            black_box(cursor.next().await?);
+            let Some(row) = cursor.next().await? else {
+                return Err(Error::NotEnoughData);
+            };
+            sum.a = sum.a.wrapping_add(row.a);
+            sum.b = sum.b.wrapping_add(row.b);
+            sum.c = sum.c.wrapping_add(row.c);
+            sum.d = sum.d.wrapping_add(row.d);
         }
 
+        black_box(sum);
         Ok(start.elapsed())
     }
 
-    let mut group = c.benchmark_group("select");
+    async fn select_bytes(client: Client, min_size: u64) -> Result<Duration> {
+        let start = Instant::now();
+        let mut cursor = client
+            .query("SELECT value FROM some")
+            .fetch_bytes("RowBinary")?;
+
+        let mut size = 0;
+        while size < min_size {
+            let buf = black_box(cursor.next().await?);
+            size += buf.unwrap().len() as u64;
+        }
+
+        Ok(start.elapsed())
+    }
+
+    let mut group = c.benchmark_group("rows");
     group.throughput(Throughput::Bytes(mem::size_of::<SomeRow>() as u64));
-    group.bench_function("no compression", |b| {
+    group.bench_function("uncompressed", |b| {
+        b.iter_custom(|iters| {
+            let client = Client::default()
+                .with_url(format!("http://{addr}"))
+                .with_compression(Compression::None);
+            runner.run(select_rows(client, iters))
+        })
+    });
+    #[cfg(feature = "lz4")]
+    group.bench_function("lz4", |b| {
+        b.iter_custom(|iters| {
+            let client = Client::default()
+                .with_url(format!("http://{addr}"))
+                .with_compression(Compression::Lz4);
+            runner.run(select_rows(client, iters))
+        })
+    });
+    group.finish();
+
+    const MIB: u64 = 1024 * 1024;
+    let mut group = c.benchmark_group("mbytes");
+    group.throughput(Throughput::Bytes(MIB));
+    group.bench_function("uncompressed", |b| {
         b.iter_custom(|iters| {
             let client = Client::default()
                 .with_url(format!("http://{addr}"))
                 .with_compression(Compression::None);
-            runner.run(run(client, iters))
+            runner.run(select_bytes(client, iters * MIB))
         })
     });
     #[cfg(feature = "lz4")]
@@ -90,7 +137,7 @@ fn select(c: &mut Criterion) {
             let client = Client::default()
                 .with_url(format!("http://{addr}"))
                 .with_compression(Compression::Lz4);
-            runner.run(run(client, iters))
+            runner.run(select_bytes(client, iters * MIB))
         })
     });
     group.finish();

examples/README.md

@@ -28,6 +28,8 @@ If something is missing, or you found a mistake in one of these examples, please
 - [custom_http_headers.rs](custom_http_headers.rs) - setting additional HTTP headers to the client, or overriding the generated ones
 - [query_id.rs](query_id.rs) - setting a specific `query_id` on the query level
 - [session_id.rs](session_id.rs) - using the client in the session context with temporary tables
+- [stream_into_file.rs](stream_into_file.rs) - streaming the query result as raw bytes into a file in an arbitrary format. Required cargo features: `futures03`.
+- [stream_arbitrary_format_rows.rs](stream_arbitrary_format_rows.rs) - streaming the query result in an arbitrary format, row by row. Required cargo features: `futures03`.
 
 ## How to run
 

examples/stream_arbitrary_format_rows.rs

@@ -0,0 +1,33 @@
+use tokio::io::AsyncBufReadExt;
+
+use clickhouse::Client;
+
+/// An example of streaming raw data in an arbitrary format leveraging the
+/// [`AsyncBufReadExt`] helpers. In this case, the format is `JSONEachRow`.
+/// Incoming data is then split into lines, and each line is deserialized into
+/// `serde_json::Value`, a dynamic representation of JSON values.
+///
+/// Similarly, it can be used with other formats such as CSV, TSV, and others
+/// that produce each row on a new line; the only difference will be in how the
+/// data is parsed. See also: https://clickhouse.com/docs/en/interfaces/formats
+///
+/// Note: `lines()` produces a new `String` for each line, so it's not the
+/// most performant way to interate over lines.
+#[tokio::main]
+async fn main() {
+    let client = Client::default().with_url("http://localhost:8123");
+    let mut lines = client
+        .query(
+            "SELECT number, hex(randomPrintableASCII(20)) AS hex_str
+             FROM system.numbers
+             LIMIT 100",
+        )
+        .fetch_bytes("JSONEachRow")
+        .unwrap()
+        .lines();
+
+    while let Some(line) = lines.next_line().await.unwrap() {
+        let value: serde_json::Value = serde_json::de::from_str(&line).unwrap();
+        println!("JSONEachRow value: {}", value);
+    }
+}

examples/stream_into_file.rs

@@ -0,0 +1,78 @@
+use clickhouse::{query::BytesCursor, Client};
+use std::time::Instant;
+use tokio::{fs::File, io::AsyncWriteExt};
+
+// Examples of streaming the result of a query in an arbitrary format into a
+// file. In this case, `CSVWithNamesAndTypes` format is used.
+// Check also other formats in https://clickhouse.com/docs/en/interfaces/formats.
+//
+// Note: there is no need to wrap `File` into `BufWriter` because `BytesCursor`
+// is buffered internally already and produces chunks of data.
+
+const NUMBERS: u32 = 100_000;
+
+fn query(numbers: u32) -> BytesCursor {
+    let client = Client::default().with_url("http://localhost:8123");
+
+    client
+        .query(
+            "SELECT number, hex(randomPrintableASCII(20)) AS hex_str
+             FROM system.numbers
+             LIMIT {limit: Int32}",
+        )
+        .param("limit", numbers)
+        .fetch_bytes("CSVWithNamesAndTypes")
+        .unwrap()
+}
+
+// Pattern 1: use the `tokio::io::copy_buf` helper.
+//
+// It shows integration with `tokio::io::AsyncBufWriteExt` trait.
+async fn tokio_copy_buf(filename: &str) {
+    let mut cursor = query(NUMBERS);
+    let mut file = File::create(filename).await.unwrap();
+    tokio::io::copy_buf(&mut cursor, &mut file).await.unwrap();
+}
+
+// Pattern 2: use `BytesCursor::next()`.
+async fn cursor_next(filename: &str) {
+    let mut cursor = query(NUMBERS);
+    let mut file = File::create(filename).await.unwrap();
+
+    while let Some(bytes) = cursor.next().await.unwrap() {
+        file.write_all(&bytes).await.unwrap();
+        println!("chunk of {}B written to {filename}", bytes.len());
+    }
+}
+
+// Pattern 3: use the `futures::(Try)StreamExt` traits.
+#[cfg(feature = "futures03")]
+async fn futures03_stream(filename: &str) {
+    use futures::TryStreamExt;
+
+    let mut cursor = query(NUMBERS);
+    let mut file = File::create(filename).await.unwrap();
+
+    while let Some(bytes) = cursor.try_next().await.unwrap() {
+        file.write_all(&bytes).await.unwrap();
+        println!("chunk of {}B written to {filename}", bytes.len());
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let start = Instant::now();
+    tokio_copy_buf("output-1.csv").await;
+    println!("written to output-1.csv in {:?}", start.elapsed());
+
+    let start = Instant::now();
+    cursor_next("output-2.csv").await;
+    println!("written to output-2.csv in {:?}", start.elapsed());
+
+    #[cfg(feature = "futures03")]
+    {
+        let start = Instant::now();
+        futures03_stream("output-3.csv").await;
+        println!("written to output-3.csv in {:?}", start.elapsed());
+    }
+}

src/bytes_ext.rs

@@ -17,6 +17,12 @@ impl BytesExt {
         self.bytes.len() - self.cursor
     }
 
+    #[inline(always)]
+    pub(crate) fn is_empty(&self) -> bool {
+        debug_assert!(self.cursor <= self.bytes.len());
+        self.cursor >= self.bytes.len()
+    }
+
     #[inline(always)]
     pub(crate) fn set_remaining(&mut self, n: usize) {
         // We can use `bytes.advance()` here, but it's slower.
@@ -26,13 +32,15 @@ impl BytesExt {
     #[cfg(any(test, feature = "lz4", feature = "watch"))]
     #[inline(always)]
     pub(crate) fn advance(&mut self, n: usize) {
+        debug_assert!(n <= self.remaining());
+
         // We can use `bytes.advance()` here, but it's slower.
         self.cursor += n;
     }
 
     #[inline(always)]
     pub(crate) fn extend(&mut self, chunk: Bytes) {
-        if self.cursor == self.bytes.len() {
+        if self.is_empty() {
             // Most of the time, we read the next chunk after consuming the previous one.
             self.bytes = chunk;
             self.cursor = 0;