add priority_map test covering "null id and value"

Author: MassivePizza

datafusion/physical-plan/src/aggregates/topk/hash_table.rs

@@ -36,11 +36,12 @@ use std::fmt::Debug;
 use std::hash::BuildHasher;
 use std::sync::Arc;
 
-/// `heap_idx` assigned to groups whose aggregate values are all NULL. Such
-/// groups are tracked in the hash table only (they never enter the heap), so
-/// they can be emitted with a NULL aggregate value at the end.
+/// `heap_idx` assigned to groups whose aggregate values are all NULL.
+/// Such groups are tracked in the hash table only, without entering heap,
+/// so they can be emitted with a NULL aggregate value at the end.
 const NULL_HEAP_IDX: usize = usize::MAX - 1;
 
+/// `heap_idx` needs to distinguish vacant slots from items that have a NULL key and value.
 const VACANT_HEAP_IDX: usize = usize::MAX;
 
 /// An entry in our hash table that:

datafusion/physical-plan/src/aggregates/topk/priority_map.rs

@@ -615,6 +615,31 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn should_emit_null_id_with_all_null_values() -> Result<()> {
+        let ids: ArrayRef = Arc::new(StringArray::from(vec![None, Some("1")]));
+        let vals: ArrayRef = Arc::new(Int64Array::from(vec![None, Some(5)]));
+        let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?;
+        agg.set_batch(ids, vals);
+        agg.insert_null(0);
+        agg.insert_with_null_groups(1)?;
+
+        let cols = agg.emit()?;
+        let batch = RecordBatch::try_new(test_schema(), cols)?;
+        let actual = format!("{}", pretty_format_batches(&[batch])?);
+        assert_snapshot!(actual, @r"
+        +----------+--------------+
+        | trace_id | timestamp_ms |
+        +----------+--------------+
+        | 1        | 5            |
+        |          |              |
+        +----------+--------------+
+        "
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn should_cap_null_groups_at_limit() -> Result<()> {
         let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5"]));