fix: prepare the nyc_taxi database ahead of benchmark (#279)

Author: timvermeulen

benches/select_nyc_taxi_data.rs

@@ -3,7 +3,7 @@
 use crate::common_select::{
     do_select_bench, print_header, print_results, BenchmarkRow, WithAccessType, WithId,
 };
-use clickhouse::{Compression, Row};
+use clickhouse::{Client, Compression, Row};
 use serde::Deserialize;
 use serde_repr::Deserialize_repr;
 use time::OffsetDateTime;
@@ -74,6 +74,86 @@ struct TripSmallMapAccess {
 impl_benchmark_row!(TripSmallSeqAccess, trip_id, "seq");
 impl_benchmark_row!(TripSmallMapAccess, trip_id, "map");
 
+// See https://clickhouse.com/docs/getting-started/example-datasets/nyc-taxi
+async fn prepare_data() {
+    let client = Client::default().with_url("http://localhost:8123");
+
+    client
+        .query("CREATE DATABASE IF NOT EXISTS nyc_taxi")
+        .execute()
+        .await
+        .unwrap();
+    client
+        .query(
+            r#"
+            CREATE TABLE IF NOT EXISTS nyc_taxi.trips_small (
+                trip_id             UInt32,
+                pickup_datetime     DateTime,
+                dropoff_datetime    DateTime,
+                pickup_longitude    Nullable(Float64),
+                pickup_latitude     Nullable(Float64),
+                dropoff_longitude   Nullable(Float64),
+                dropoff_latitude    Nullable(Float64),
+                passenger_count     UInt8,
+                trip_distance       Float32,
+                fare_amount         Float32,
+                extra               Float32,
+                tip_amount          Float32,
+                tolls_amount        Float32,
+                total_amount        Float32,
+                payment_type        Enum('CSH' = 1, 'CRE' = 2, 'NOC' = 3, 'DIS' = 4, 'UNK' = 5),
+                pickup_ntaname      LowCardinality(String),
+                dropoff_ntaname     LowCardinality(String)
+            )
+            ENGINE = MergeTree
+            PRIMARY KEY (pickup_datetime, dropoff_datetime)
+            "#,
+        )
+        .execute()
+        .await
+        .unwrap();
+
+    let len = client
+        .query("SELECT count() FROM nyc_taxi.trips_small")
+        .fetch_one::<usize>()
+        .await
+        .unwrap();
+
+    if len == 0 {
+        client
+        .query(
+            "
+            INSERT INTO nyc_taxi.trips_small
+            SELECT
+                trip_id,
+                pickup_datetime,
+                dropoff_datetime,
+                pickup_longitude,
+                pickup_latitude,
+                dropoff_longitude,
+                dropoff_latitude,
+                passenger_count,
+                trip_distance,
+                fare_amount,
+                extra,
+                tip_amount,
+                tolls_amount,
+                total_amount,
+                payment_type,
+                pickup_ntaname,
+                dropoff_ntaname
+            FROM gcs(
+                'https://storage.googleapis.com/clickhouse-public-datasets/nyc-taxi/trips_{0..2}.gz',
+                'TabSeparatedWithNames'
+            );
+            ",
+        )
+        .execute()
+        .await
+        .unwrap();
+    }
+}
+
 async fn bench<T: BenchmarkRow>(compression: Compression, validation: bool) {
     let stats = do_select_bench::<T>(
         "SELECT * FROM nyc_taxi.trips_small ORDER BY trip_id DESC",
@@ -87,6 +167,7 @@ async fn bench<T: BenchmarkRow>(compression: Compression, validation: bool) {
 
 #[tokio::main]
 async fn main() {
+    prepare_data().await;
     print_header(Some("  access"));
     bench::<TripSmallSeqAccess>(Compression::None, false).await;
     bench::<TripSmallSeqAccess>(Compression::None, true).await;