Wrong results when grouping with dictionary arrays with nulls

[alamb]

Describe the bug
Nulls are ignored when aggregating dictionary arrays.

To Reproduce
Create a dictionary array with nulls and then aggregate using that dictionary

Here is an example:

async fn dictionary_test() {
    let key_builder = PrimitiveBuilder::<Int8Type>::new(100);
    let value_builder = StringBuilder::new(100);
    let mut builder = StringDictionaryBuilder::new(key_builder, value_builder);

    builder.append("b").unwrap();
    builder.append_null().unwrap();
    builder.append("b").unwrap();
    builder.append("a").unwrap();
    builder.append("c").unwrap();
    builder.append_null().unwrap();

    let array = builder.finish();
    print_array(&array);

    // Now, let's run a query showing the invalid values are being used as well:
    let batch = RecordBatch::try_from_iter(vec![("dict", array)]).unwrap();
    let table = MemTable::try_new(batch.schema(), vec![vec![batch]]).unwrap();
    let table = Arc::new(table);

    // +------+
    // | dict |
    // +------+
    // | b    |
    // |      |
    // | b    |
    // | a    |
    // | c    |
    // |      |
    // +------+
    run_query(table.clone(), "select * from t").await;

    // now aggregate the query
    // (wrong results -- there are only 2 values for b, but it shows up as 4)
    // +------+-----------------+
    // | dict | COUNT(UInt8(1)) |
    // +------+-----------------+
    // | a    | 1               |
    // | c    | 1               |
    // | b    | 4               |
    // +------+-----------------+
    run_query(table.clone(), "select dict, count(*) from t group by dict").await;

}

Here is a simply reproducer in sql.rs:

diff --git a/datafusion/tests/sql.rs b/datafusion/tests/sql.rs
index d9f7c6ea4..a323fc499 100644
--- a/datafusion/tests/sql.rs
+++ b/datafusion/tests/sql.rs
@@ -3019,6 +3019,15 @@ async fn query_on_string_dictionary() -> Result<()> {
     let expected = vec![vec!["2"]];
     assert_eq!(expected, actual);
 
+    // grouping
+    let sql = "SELECT d1, COUNT(*) FROM test group by d1";
+    let actual = execute(&mut ctx, sql).await;
+    let expected = vec![
+        vec!["one", "1"],
+        vec!["three", "1"],
+    ];
+    assert_eq!(expected, actual);
+
     Ok(())
 }

results in

---- query_on_string_dictionary stdout ----
thread 'query_on_string_dictionary' panicked at 'assertion failed: `(left == right)`
  left: `[["one", "1"], ["three", "1"]]`,
 right: `[["one", "2"], ["three", "1"]]`', datafusion/tests/sql.rs:3029:5
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

Expected behavior
Null entries should be ignored, but instead whatever value in the keys array is at that position is used to find a value.

Additional context
We found this in IOx (where @tustvold clevery used -1 for entries in keys that were NULL rather than 0): https://github.com/influxdata/influxdb_iox/issues/2112

[jorgecarleitao]

This may be related to the fact that both the dict and the values are nullable. I do not know how the builder is populating the validity, but there is both array.keys().nulls() and array.values().nulls() around.

[alamb]

This may be related to the fact that both the dict and the values are nullable. I do not know how the builder is populating the validity, but there is both array.keys().nulls() and array.values().nulls() around.

Thanks @jorgecarleitao !

In my investigation, I believe the issue is that the values in array.keys() are marked as invalid, but the code that unpacks dictonaries into ScalarValues (which I wrote 🤦 ), for example https://github.com/apache/arrow-datafusion/blob/master/datafusion/src/physical_plan/hash_aggregate.rs#L504 does not check is_valid (and there are at least three such places). I am digging in more now

The keys
array[0]: 0, valid(true)
array[1]: 0, valid(false)
array[2]: 0, valid(true)
array[3]: 1, valid(true)
array[4]: 2, valid(true)
array[5]: 0, valid(false)
The values
array[0]: "b", valid(true)
array[1]: "a", valid(true)
array[2]: "c", valid(true)
The array:
array[0]: "b", valid(true)
array[1]: "", valid(false)
array[2]: "b", valid(true)
array[3]: "a", valid(true)
array[4]: "c", valid(true)
array[5]: "", valid(false)

[alamb]

BTW when grouping, according to postgres, null should be included:

alamb=# select * from test;
  d1   
-------
 one
 
 three
(3 rows)

alamb=# select d1, count(*) from test group by d1;
  d1   | count 
-------+-------
       |     1
 three |     1
 one   |     1
(3 rows)

(note the column for d1 = NULL)

[alamb]

So looking more carefully I think this is a more general problem related to grouping on NULLs: #782 (the group by code does not check for null)