Merge branch 'main' into alamb/update_arrow_58

Author: alamb

Cargo.lock

@@ -183,7 +183,7 @@ regex = "1.12"
 rstest = "0.26.1"
 serde_json = "1"
 sha2 = "^0.10.9"
-sqlparser = { version = "0.60.0", default-features = false, features = ["std", "visitor"] }
+sqlparser = { version = "0.61.0", default-features = false, features = ["std", "visitor"] }
 strum = "0.27.2"
 strum_macros = "0.27.2"
 tempfile = "3"

Cargo.toml

@@ -44,7 +44,7 @@ fn make_settings() -> Settings {
     settings
 }
 
-async fn setup_minio_container() -> ContainerAsync<minio::MinIO> {
+async fn setup_minio_container() -> Result<ContainerAsync<minio::MinIO>, String> {
     const MINIO_ROOT_USER: &str = "TEST-DataFusionLogin";
     const MINIO_ROOT_PASSWORD: &str = "TEST-DataFusionPassword";
 
@@ -99,27 +99,23 @@ async fn setup_minio_container() -> ContainerAsync<minio::MinIO> {
                     let stdout = container.stdout_to_vec().await.unwrap_or_default();
                     let stderr = container.stderr_to_vec().await.unwrap_or_default();
 
-                    panic!(
+                    return Err(format!(
                         "Failed to execute command: {}\nError: {}\nStdout: {:?}\nStderr: {:?}",
                         cmd_ref,
                         e,
                         String::from_utf8_lossy(&stdout),
                         String::from_utf8_lossy(&stderr)
-                    );
+                    ));
                 }
             }
 
-            container
+            Ok(container)
         }
 
-        Err(TestcontainersError::Client(e)) => {
-            panic!(
-                "Failed to start MinIO container. Ensure Docker is running and accessible: {e}"
-            );
-        }
-        Err(e) => {
-            panic!("Failed to start MinIO container: {e}");
-        }
+        Err(TestcontainersError::Client(e)) => Err(format!(
+            "Failed to start MinIO container. Ensure Docker is running and accessible: {e}"
+        )),
+        Err(e) => Err(format!("Failed to start MinIO container: {e}")),
     }
 }
 
@@ -253,7 +249,14 @@ async fn test_cli() {
         return;
     }
 
-    let container = setup_minio_container().await;
+    let container = match setup_minio_container().await {
+        Ok(c) => c,
+        Err(e) if e.contains("toomanyrequests") => {
+            eprintln!("Skipping test: Docker pull rate limit reached: {e}");
+            return;
+        }
+        e @ Err(_) => e.unwrap(),
+    };
 
     let settings = make_settings();
     let _bound = settings.bind_to_scope();
@@ -286,7 +289,14 @@ async fn test_aws_options() {
     let settings = make_settings();
     let _bound = settings.bind_to_scope();
 
-    let container = setup_minio_container().await;
+    let container = match setup_minio_container().await {
+        Ok(c) => c,
+        Err(e) if e.contains("toomanyrequests") => {
+            eprintln!("Skipping test: Docker pull rate limit reached: {e}");
+            return;
+        }
+        e @ Err(_) => e.unwrap(),
+    };
     let port = container.get_host_port_ipv4(9000).await.unwrap();
 
     let input = format!(
@@ -377,7 +387,14 @@ async fn test_s3_url_fallback() {
         return;
     }
 
-    let container = setup_minio_container().await;
+    let container = match setup_minio_container().await {
+        Ok(c) => c,
+        Err(e) if e.contains("toomanyrequests") => {
+            eprintln!("Skipping test: Docker pull rate limit reached: {e}");
+            return;
+        }
+        e @ Err(_) => e.unwrap(),
+    };
 
     let mut settings = make_settings();
     settings.set_snapshot_suffix("s3_url_fallback");
@@ -407,8 +424,14 @@ async fn test_object_store_profiling() {
         return;
     }
 
-    let container = setup_minio_container().await;
-
+    let container = match setup_minio_container().await {
+        Ok(c) => c,
+        Err(e) if e.contains("toomanyrequests") => {
+            eprintln!("Skipping test: Docker pull rate limit reached: {e}");
+            return;
+        }
+        e @ Err(_) => e.unwrap(),
+    };
     let mut settings = make_settings();
 
     // as the object store profiling contains timestamps and durations, we must

datafusion-cli/tests/cli_integration.rs

@@ -251,6 +251,39 @@ fn criterion_benchmark(c: &mut Criterion) {
             )
         })
     });
+
+    c.bench_function("array_agg_query_group_by_few_groups", |b| {
+        b.iter(|| {
+            query(
+                ctx.clone(),
+                &rt,
+                "SELECT u64_narrow, array_agg(f64) \
+                 FROM t GROUP BY u64_narrow",
+            )
+        })
+    });
+
+    c.bench_function("array_agg_query_group_by_mid_groups", |b| {
+        b.iter(|| {
+            query(
+                ctx.clone(),
+                &rt,
+                "SELECT u64_mid, array_agg(f64) \
+                 FROM t GROUP BY u64_mid",
+            )
+        })
+    });
+
+    c.bench_function("array_agg_query_group_by_many_groups", |b| {
+        b.iter(|| {
+            query(
+                ctx.clone(),
+                &rt,
+                "SELECT u64_wide, array_agg(f64) \
+                 FROM t GROUP BY u64_wide",
+            )
+        })
+    });
 }
 
 criterion_group!(benches, criterion_benchmark);

datafusion/core/benches/aggregate_query_sql.rs

@@ -45,7 +45,7 @@ pub fn create_table_provider(
 ) -> Result<Arc<MemTable>> {
     let schema = Arc::new(create_schema());
     let partitions =
-        create_record_batches(schema.clone(), array_len, partitions_len, batch_size);
+        create_record_batches(&schema, array_len, partitions_len, batch_size);
     // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
     MemTable::try_new(schema, partitions).map(Arc::new)
 }
@@ -56,21 +56,19 @@ pub fn create_schema() -> Schema {
         Field::new("utf8", DataType::Utf8, false),
         Field::new("f32", DataType::Float32, false),
         Field::new("f64", DataType::Float64, true),
-        // This field will contain integers randomly selected from a large
-        // range of values, i.e. [0, u64::MAX], such that there are none (or
-        // very few) repeated values.
-        Field::new("u64_wide", DataType::UInt64, true),
-        // This field will contain integers randomly selected from a narrow
-        // range of values such that there are a few distinct values, but they
-        // are repeated often.
+        // Integers randomly selected from a wide range of values, i.e. [0,
+        // u64::MAX], such that there are ~no repeated values.
+        Field::new("u64_wide", DataType::UInt64, false),
+        // Integers randomly selected from a mid-range of values [0, 1000),
+        // providing ~1000 distinct groups.
+        Field::new("u64_mid", DataType::UInt64, false),
+        // Integers randomly selected from a narrow range of values such that
+        // there are a few distinct values, but they are repeated often.
         Field::new("u64_narrow", DataType::UInt64, false),
     ])
 }
 
-fn create_data(size: usize, null_density: f64) -> Vec<Option<f64>> {
-    // use random numbers to avoid spurious compiler optimizations wrt to branching
-    let mut rng = StdRng::seed_from_u64(42);
-
+fn create_data(rng: &mut StdRng, size: usize, null_density: f64) -> Vec<Option<f64>> {
     (0..size)
         .map(|_| {
             if rng.random::<f64>() > null_density {
@@ -82,56 +80,43 @@ fn create_data(size: usize, null_density: f64) -> Vec<Option<f64>> {
         .collect()
 }
 
-fn create_integer_data(
-    rng: &mut StdRng,
-    size: usize,
-    value_density: f64,
-) -> Vec<Option<u64>> {
-    (0..size)
-        .map(|_| {
-            if rng.random::<f64>() > value_density {
-                None
-            } else {
-                Some(rng.random::<u64>())
-            }
-        })
-        .collect()
-}
-
 fn create_record_batch(
     schema: SchemaRef,
     rng: &mut StdRng,
     batch_size: usize,
-    i: usize,
+    batch_index: usize,
 ) -> RecordBatch {
-    // the 4 here is the number of different keys.
-    // a higher number increase sparseness
-    let vs = [0, 1, 2, 3];
-    let keys: Vec<String> = (0..batch_size)
-        .map(
-            // use random numbers to avoid spurious compiler optimizations wrt to branching
-            |_| format!("hi{:?}", vs.choose(rng)),
-        )
-        .collect();
-    let keys: Vec<&str> = keys.iter().map(|e| &**e).collect();
+    // Randomly choose from 4 distinct key values; a higher number increases sparseness.
+    let key_suffixes = [0, 1, 2, 3];
+    let keys = StringArray::from_iter_values(
+        (0..batch_size).map(|_| format!("hi{}", key_suffixes.choose(rng).unwrap())),
+    );
 
-    let values = create_data(batch_size, 0.5);
+    let values = create_data(rng, batch_size, 0.5);
 
     // Integer values between [0, u64::MAX].
-    let integer_values_wide = create_integer_data(rng, batch_size, 9.0);
+    let integer_values_wide = (0..batch_size)
+        .map(|_| rng.random::<u64>())
+        .collect::<Vec<_>>();
 
-    // Integer values between [0, 9].
+    // Integer values between [0, 1000).
+    let integer_values_mid = (0..batch_size)
+        .map(|_| rng.random_range(0..1000))
+        .collect::<Vec<_>>();
+
+    // Integer values between [0, 10).
     let integer_values_narrow = (0..batch_size)
-        .map(|_| rng.random_range(0_u64..10))
+        .map(|_| rng.random_range(0..10))
         .collect::<Vec<_>>();
 
     RecordBatch::try_new(
         schema,
         vec![
-            Arc::new(StringArray::from(keys)),
-            Arc::new(Float32Array::from(vec![i as f32; batch_size])),
+            Arc::new(keys),
+            Arc::new(Float32Array::from(vec![batch_index as f32; batch_size])),
             Arc::new(Float64Array::from(values)),
             Arc::new(UInt64Array::from(integer_values_wide)),
+            Arc::new(UInt64Array::from(integer_values_mid)),
             Arc::new(UInt64Array::from(integer_values_narrow)),
         ],
     )
@@ -140,21 +125,29 @@ fn create_record_batch(
 
 /// Create record batches of `partitions_len` partitions and `batch_size` for each batch,
 /// with a total number of `array_len` records
-#[expect(clippy::needless_pass_by_value)]
 pub fn create_record_batches(
-    schema: SchemaRef,
+    schema: &SchemaRef,
     array_len: usize,
     partitions_len: usize,
     batch_size: usize,
 ) -> Vec<Vec<RecordBatch>> {
     let mut rng = StdRng::seed_from_u64(42);
-    (0..partitions_len)
-        .map(|_| {
-            (0..array_len / batch_size / partitions_len)
-                .map(|i| create_record_batch(schema.clone(), &mut rng, batch_size, i))
-                .collect::<Vec<_>>()
-        })
-        .collect::<Vec<_>>()
+    let mut partitions = Vec::with_capacity(partitions_len);
+    let batches_per_partition = array_len / batch_size / partitions_len;
+
+    for _ in 0..partitions_len {
+        let mut batches = Vec::with_capacity(batches_per_partition);
+        for batch_index in 0..batches_per_partition {
+            batches.push(create_record_batch(
+                schema.clone(),
+                &mut rng,
+                batch_size,
+                batch_index,
+            ));
+        }
+        partitions.push(batches);
+    }
+    partitions
 }
 
 /// An enum that wraps either a regular StringBuilder or a GenericByteViewBuilder