Fix bug that COUNT(DISTINCT) on StringView panics (#11768)

* fix bug * Add test showing panic on string view --------- Co-authored-by: Andrew Lamb <[email protected]>
apache · Aug 1, 2024 · f044bc8 · f044bc8
1 parent 0d98b99
commit f044bc8
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 1 deletion.
diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs
@@ -237,14 +237,17 @@ impl AggregateUDFImpl for Count {
                 Box::new(BytesDistinctCountAccumulator::<i32>::new(OutputType::Utf8))
             }
             DataType::Utf8View => {
-                Box::new(BytesViewDistinctCountAccumulator::new(OutputType::Utf8))
+                Box::new(BytesViewDistinctCountAccumulator::new(OutputType::Utf8View))
             }
             DataType::LargeUtf8 => {
                 Box::new(BytesDistinctCountAccumulator::<i64>::new(OutputType::Utf8))
             }
             DataType::Binary => Box::new(BytesDistinctCountAccumulator::<i32>::new(
                 OutputType::Binary,
             )),
+            DataType::BinaryView => Box::new(BytesViewDistinctCountAccumulator::new(
+                OutputType::BinaryView,
+            )),
             DataType::LargeBinary => Box::new(BytesDistinctCountAccumulator::<i64>::new(
                 OutputType::Binary,
             )),

diff --git a/datafusion/sqllogictest/test_files/string_view.slt b/datafusion/sqllogictest/test_files/string_view.slt
@@ -321,6 +321,40 @@ logical_plan
 02)--Filter: CAST(test.column2_utf8 AS Utf8View) = test.column1_utf8view
 03)----TableScan: test projection=[column1_utf8, column2_utf8, column1_utf8view]
 
+## Test distinct aggregates
+query III
+SELECT
+  COUNT(DISTINCT column1_utf8),
+  COUNT(DISTINCT column1_utf8view),
+  COUNT(DISTINCT column1_dict)
+FROM test;
+----
+3 3 3
+
+query III
+SELECT
+  COUNT(DISTINCT column1_utf8),
+  COUNT(DISTINCT column1_utf8view),
+  COUNT(DISTINCT column1_dict)
+FROM test
+GROUP BY column2_utf8view;
+----
+1 1 1
+1 1 1
+1 1 1
+
+
+query TT
+EXPLAIN SELECT
+  COUNT(DISTINCT column1_utf8),
+  COUNT(DISTINCT column1_utf8view),
+  COUNT(DISTINCT column1_dict)
+FROM test;
+----
+logical_plan
+01)Aggregate: groupBy=[[]], aggr=[[count(DISTINCT test.column1_utf8), count(DISTINCT test.column1_utf8view), count(DISTINCT test.column1_dict)]]
+02)--TableScan: test projection=[column1_utf8, column1_utf8view, column1_dict]
+
 
 statement ok
 drop table test;