jaylmiller commented on code in PR #5554:
URL: https://github.com/apache/arrow-datafusion/pull/5554#discussion_r1133280472
##########
datafusion/physical-expr/src/aggregate/count_distinct.rs:
##########
@@ -577,4 +746,106 @@ mod tests {
assert_eq!(result, ScalarValue::Int64(Some(2)));
Ok(())
}
+
+ #[test]
+ fn count_distinct_dict_update() -> Result<()> {
+ let values = StringArray::from_iter_values(["a", "b", "c"]);
+ // value "b" is never used
+ let keys =
+ Int8Array::from_iter(vec![Some(0), Some(0), Some(0), Some(0),
None, Some(2)]);
+ let arrays =
+ vec![
+ Arc::new(DictionaryArray::<Int8Type>::try_new(&keys,
&values).unwrap())
+ as ArrayRef,
+ ];
+ let agg = DistinctCount::new(
+ arrays[0].data_type().clone(),
+ Arc::new(NoOp::new()),
+ String::from("__col_name__"),
+ );
+ let mut accum = agg.create_accumulator()?;
+ accum.update_batch(&arrays)?;
+ // should evaluate to 2 since "b" never seen
+ assert_eq!(accum.evaluate()?, ScalarValue::Int64(Some(2)));
+ // now update with a new batch that does use "b"
+ let values = StringArray::from_iter_values(["a", "b", "c"]);
+ let keys = Int8Array::from_iter(vec![Some(1), Some(1), None]);
+ let arrays =
+ vec![
+ Arc::new(DictionaryArray::<Int8Type>::try_new(&keys,
&values).unwrap())
+ as ArrayRef,
+ ];
+ accum.update_batch(&arrays)?;
+ assert_eq!(accum.evaluate()?, ScalarValue::Int64(Some(3)));
+ Ok(())
+ }
+
+ #[test]
+ fn count_distinct_dict_merge() -> Result<()> {
+ let values = StringArray::from_iter_values(["a", "b", "c"]);
+ let keys = Int8Array::from_iter(vec![Some(0), Some(0), None]);
+ let arrays =
+ vec![
+ Arc::new(DictionaryArray::<Int8Type>::try_new(&keys,
&values).unwrap())
+ as ArrayRef,
+ ];
+ let agg = DistinctCount::new(
+ arrays[0].data_type().clone(),
+ Arc::new(NoOp::new()),
+ String::from("__col_name__"),
+ );
+ // create accum with 1 value seen
+ let mut accum = agg.create_accumulator()?;
+ accum.update_batch(&arrays)?;
+ assert_eq!(accum.evaluate()?, ScalarValue::Int64(Some(1)));
+ // create accum with state that has seen "a" and "b" but not "c"
+ let values = StringArray::from_iter_values(["a", "b", "c"]);
+ let keys = Int8Array::from_iter(vec![Some(0), Some(1), None]);
Review Comment:
Yes this would break under my current assumption, will change some stuff and
get this working again.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
