ARROW-9937: [Rust] [DataFusion] Improved aggregations

rust/arrow/src/compute/kernels/aggregate.rs

@@ -19,9 +19,44 @@
 
 use std::ops::Add;
 
-use crate::array::{Array, PrimitiveArray};
+use crate::array::{
+ Array, LargeStringArray, PrimitiveArray, StringArray, StringArrayOps,
+};
 use crate::datatypes::ArrowNumericType;
 
+/// Helper macro to perform min/max of strings
+macro_rules! min_max_string_helper {
+ ($array:expr, $cmp:tt) => {{
+ let null_count = $array.null_count();
+
+ if null_count == $array.len() {
+ return None
+ }
+ let mut n = "";
+ let mut has_value = false;
+ let data = $array.data();
+
+ if null_count == 0 {
+ for i in 0..data.len() {
+ let item = $array.value(i);
+ if !has_value || (&n $cmp &item) {
+ has_value = true;
+ n = item;
+ }
+ }
+ } else {
+ for i in 0..data.len() {
+ let item = $array.value(i);
+ if !has_value || data.is_valid(i) && (&n $cmp &item) {
+ has_value = true;
+ n = item;
+ }
+ }
+ }
+ Some(n)
+ }}
+}
+
 /// Returns the minimum value in the array, according to the natural order.
 pub fn min<T>(array: &PrimitiveArray<T>) -> Option<T::Native>
 where
@@ -38,6 +73,26 @@ where
  min_max_helper(array, |a, b| a < b)
 }
 
+/// Returns the maximum value in the string array, according to the natural order.
+pub fn max_string(array: &StringArray) -> Option<&str> {
+ min_max_string_helper!(array, <)
+}
+
+/// Returns the minimum value in the string array, according to the natural order.
+pub fn min_string(array: &StringArray) -> Option<&str> {
+ min_max_string_helper!(array, >)
+}
+
+/// Returns the minimum value in the string array, according to the natural order.
+pub fn max_large_string(array: &LargeStringArray) -> Option<&str> {
+ min_max_string_helper!(array, <)
+}
+
+/// Returns the minimum value in the string array, according to the natural order.
+pub fn min_large_string(array: &LargeStringArray) -> Option<&str> {
+ min_max_string_helper!(array, >)
+}
+
 /// Helper function to perform min/max lambda function on values from a numeric array.
 fn min_max_helper<T, F>(array: &PrimitiveArray<T>, cmp: F) -> Option<T::Native>
 where
@@ -149,4 +204,18 @@ mod tests {
  assert_eq!(5, min(&a).unwrap());
  assert_eq!(9, max(&a).unwrap());
  }
+
+ #[test]
+ fn test_string_min_max_with_nulls() {
+ let a = StringArray::from(vec![Some("b"), None, None, Some("a"), Some("c")]);
+ assert_eq!("a", min_string(&a).unwrap());
+ assert_eq!("c", max_string(&a).unwrap());
+ }
+
+ #[test]
+ fn test_string_min_max_all_nulls() {
+ let a = StringArray::from(vec![None, None]);
+ assert_eq!(None, min_string(&a));
+ assert_eq!(None, max_string(&a));
+ }
 }

rust/datafusion/Cargo.toml

@@ -56,6 +56,7 @@ num_cpus = "1.13.0"
 chrono = "0.4"
 
 [dev-dependencies]
+rand = "0.7"
 criterion = "0.3"
 tempfile = "3"
 futures = "0.3"

rust/datafusion/benches/aggregate_query_sql.rs

@@ -19,18 +19,26 @@
 extern crate criterion;
 use criterion::Criterion;
 
-use std::env;
+use rand::seq::SliceRandom;
+use rand::Rng;
 use std::sync::Arc;
 
 extern crate arrow;
 extern crate datafusion;
 
-use arrow::datatypes::{DataType, Field, Schema};
+use arrow::{
+ array::Float32Array,
+ array::Float64Array,
+ array::StringArray,
+ datatypes::{DataType, Field, Schema},
+ record_batch::RecordBatch,
+};
 
-use datafusion::datasource::{CsvFile, CsvReadOptions, MemTable};
+use datafusion::datasource::MemTable;
+use datafusion::error::Result;
 use datafusion::execution::context::ExecutionContext;
 
-fn aggregate_query(ctx: &mut ExecutionContext, sql: &str) {
+fn query(ctx: &mut ExecutionContext, sql: &str) {
  // execute the query
  let df = ctx.sql(&sql).unwrap();
  let results = df.collect().unwrap();
@@ -39,72 +47,107 @@ fn aggregate_query(ctx: &mut ExecutionContext, sql: &str) {
  for _batch in results {}
 }
 
-fn create_context() -> ExecutionContext {
- // define schema for data source (csv file)
+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 = rand::thread_rng();
+
+ (0..size)
+ .map(|_| {
+ if rng.gen::<f64>() > null_density {
+ None
+ } else {
+ Some(rng.gen::<f64>())
+ }
+ })
+ .collect()
+}
+
+fn create_context(
+ partitions_len: usize,
+ array_len: usize,
+ batch_size: usize,
+) -> Result<ExecutionContext> {
+ // define a schema.
  let schema = Arc::new(Schema::new(vec![
- Field::new("c1", DataType::Utf8, false),
- Field::new("c2", DataType::UInt32, false),
- Field::new("c3", DataType::Int8, false),
- Field::new("c4", DataType::Int16, false),
- Field::new("c5", DataType::Int32, false),
- Field::new("c6", DataType::Int64, false),
- Field::new("c7", DataType::UInt8, false),
- Field::new("c8", DataType::UInt16, false),
- Field::new("c9", DataType::UInt32, false),
- Field::new("c10", DataType::UInt64, false),
- Field::new("c11", DataType::Float32, false),
- Field::new("c12", DataType::Float64, false),
- Field::new("c13", DataType::Utf8, false),
+ Field::new("utf8", DataType::Utf8, false),
+ Field::new("f32", DataType::Float32, false),
+ Field::new("f64", DataType::Float64, false),
  ]));
 
- let testdata = env::var("ARROW_TEST_DATA").expect("ARROW_TEST_DATA not defined");
+ // define data.
+ let partitions = (0..partitions_len)
+ .map(|_| {
+ (0..array_len / batch_size / partitions_len)
+ .map(|i| {
+ // the 4 here is the number of different keys.
+ // a higher number increase sparseness
+ let vs = vec![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(&mut rand::thread_rng())),
+ )
+ .collect();
+ let keys: Vec<&str> = keys.iter().map(|e| &**e).collect();
+
+ let values = create_data(batch_size, 0.5);
+
+ RecordBatch::try_new(
+ schema.clone(),
+ vec![
+ Arc::new(StringArray::from(keys)),
+ Arc::new(Float32Array::from(vec![i as f32; batch_size])),
+ Arc::new(Float64Array::from(values)),
+ ],
+ )
+ .unwrap()
+ })
+ .collect::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
 
- // create CSV data source
- let csv = CsvFile::try_new(
- &format!("{}/csv/aggregate_test_100.csv", testdata),
- CsvReadOptions::new().schema(&schema),
- )
- .unwrap();
+ let mut ctx = ExecutionContext::new();
 
- let mem_table = MemTable::load(&csv).unwrap();
+ // declare a table in memory. In spark API, this corresponds to createDataFrame(...).
+ let provider = MemTable::new(schema, partitions)?;
+ ctx.register_table("t", Box::new(provider));
 
- // create local execution context
- let mut ctx = ExecutionContext::new();
- ctx.register_table("aggregate_test_100", Box::new(mem_table));
- ctx
+ Ok(ctx)
 }
 
 fn criterion_benchmark(c: &mut Criterion) {
- c.bench_function("aggregate_query_no_group_by", |b| {
- let mut ctx = create_context();
+ let partitions_len = 4;
+ let array_len = 32768; // 2^15
+ let batch_size = 2048; // 2^11
+ let mut ctx = create_context(partitions_len, array_len, batch_size).unwrap();
+
+ c.bench_function("aggregate_query_no_group_by 15 12", |b| {
  b.iter(|| {
-aggregate_query(
+query(
  &mut ctx,
- "SELECT MIN(c12), MAX(c12) \
- FROM aggregate_test_100",
+ "SELECT MIN(f64), AVG(f64), COUNT(f64) \
+ FROM t",
  )
  })
  });
 
- c.bench_function("aggregate_query_group_by", |b| {
- let mut ctx = create_context();
+ c.bench_function("aggregate_query_group_by 15 12", |b| {
  b.iter(|| {
-aggregate_query(
+query(
  &mut ctx,
- "SELECT c1, MIN(c12), MAX(c12) \
- FROM aggregate_test_100 GROUP BY c1",
+ "SELECT utf8, MIN(f64), AVG(f64), COUNT(f64) \
+ FROM t GROUP BY utf8",
  )
  })
  });
 
- c.bench_function("aggregate_query_group_by_with_filter", |b| {
- let mut ctx = create_context();
+ c.bench_function("aggregate_query_group_by_with_filter 15 12", |b| {
  b.iter(|| {
-aggregate_query(
+query(
  &mut ctx,
- "SELECT c1, MIN(c12), MAX(c12) \
- FROM aggregate_test_100 \
- WHERE c11 > 0.1 AND c11 < 0.9 GROUP BY c1",
+ "SELECT utf8, MIN(f64), AVG(f64), COUNT(f64) \
+ FROM t \
+ WHERE f32 > 10 AND f32 < 20 GROUP BY utf8",
  )
  })
  });

rust/datafusion/src/execution/context.rs

@@ -516,7 +516,8 @@ mod tests {
  use crate::test;
  use crate::variable::VarType;
  use arrow::array::{
- ArrayRef, Int32Array, PrimitiveArrayOps, StringArray, StringArrayOps,
+ ArrayRef, Float64Array, Int32Array, PrimitiveArrayOps, StringArray,
+ StringArrayOps,
  };
  use arrow::compute::add;
  use std::fs::File;
@@ -847,6 +848,24 @@ mod tests {
  Ok(())
  }
 
+ #[test]
+ fn aggregate_grouped_empty() -> Result<()> {
+ let results =
+ execute("SELECT c1, AVG(c2) FROM test WHERE c1 = 123 GROUP BY c1", 4)?;
+ assert_eq!(results.len(), 1);
+
+ let batch = &results[0];
+
+ assert_eq!(field_names(batch), vec!["c1", "AVG(c2)"]);
+
+ let expected: Vec<&str> = vec![];
+ let mut rows = test::format_batch(&batch);
+ rows.sort();
+ assert_eq!(rows, expected);
+
+ Ok(())
+ }
+
  #[test]
  fn aggregate_grouped_max() -> Result<()> {
  let results = execute("SELECT c1, MAX(c2) FROM test GROUP BY c1", 4)?;
@@ -1146,6 +1165,41 @@ mod tests {
  Ok(())
  }
 
+ #[test]
+ fn simple_avg() -> Result<()> {
+ let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
+
+ let batch1 = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
+ )?;
+ let batch2 = RecordBatch::try_new(
+ Arc::new(schema.clone()),
+ vec![Arc::new(Int32Array::from(vec![4, 5]))],
+ )?;
+
+ let mut ctx = ExecutionContext::new();
+
+ let provider = MemTable::new(Arc::new(schema), vec![vec![batch1], vec![batch2]])?;
+ ctx.register_table("t", Box::new(provider));
+
+ let result = collect(&mut ctx, "SELECT AVG(a) FROM t")?;
+
+ let batch = &result[0];
+ assert_eq!(1, batch.num_columns());
+ assert_eq!(1, batch.num_rows());
+
+ let values = batch
+ .column(0)
+ .as_any()
+ .downcast_ref::<Float64Array>()
+ .expect("failed to cast version");
+ assert_eq!(values.len(), 1);
+ // avg(1,2,3,4,5) = 3.0
+ assert_eq!(values.value(0), 3.0_f64);
+ Ok(())
+ }
+
  #[test]
  fn custom_query_planner() -> Result<()> {
  let mut ctx = ExecutionContext::with_config(

rust/datafusion/src/lib.rs

@@ -66,6 +66,7 @@ pub mod logical_plan;
 pub mod optimizer;
 pub mod physical_plan;
 pub mod prelude;
+pub mod scalar;
 pub mod sql;
 pub mod variable;