Fix: group items in projection should wrap nullable too.

Author: RinChanNOWWW

src/query/sql/src/executor/physical_plan.rs

@@ -174,24 +174,29 @@ pub struct AggregateExpand {
     pub plan_id: u32,
 
     pub input: Box<PhysicalPlan>,
-    pub group_bys: Vec<usize>,
+    pub group_bys: Vec<IndexType>,
     pub grouping_id_index: IndexType,
-    pub grouping_sets: Vec<Vec<usize>>,
+    pub grouping_sets: Vec<Vec<IndexType>>,
     /// Only used for explain
     pub stat_info: Option<PlanStatsInfo>,
 }
 
 impl AggregateExpand {
     pub fn output_schema(&self) -> Result<DataSchemaRef> {
         let input_schema = self.input.output_schema()?;
-        let input_fields = input_schema.fields();
-        let mut output_fields = Vec::with_capacity(input_fields.len() + 1);
-        for field in input_fields {
-            output_fields.push(DataField::new(
-                field.name(),
-                field.data_type().wrap_nullable(),
-            ));
+        let mut output_fields = input_schema.fields().clone();
+
+        for group_by in self
+            .group_bys
+            .iter()
+            .filter(|&index| *index != self.grouping_id_index)
+        {
+            // All group by columns will wrap nullable.
+            let i = input_schema.index_of(&group_by.to_string())?;
+            let f = &mut output_fields[i];
+            *f = DataField::new(f.name(), f.data_type().wrap_nullable())
         }
+
         output_fields.push(DataField::new(
             &self.grouping_id_index.to_string(),
             DataType::Number(NumberDataType::UInt32),

src/query/sql/src/planner/semantic/grouping_check.rs

@@ -49,7 +49,8 @@ impl<'a> GroupingChecker<'a> {
             .get(&format!("{:?}", scalar))
         {
             let column = &self.bind_context.aggregate_info.group_items[*index];
-            let column_binding = if let ScalarExpr::BoundColumnRef(column_ref) = &column.scalar {
+            let mut column_binding = if let ScalarExpr::BoundColumnRef(column_ref) = &column.scalar
+            {
                 column_ref.column.clone()
             } else {
                 ColumnBinding {
@@ -61,6 +62,13 @@ impl<'a> GroupingChecker<'a> {
                     visibility: Visibility::Visible,
                 }
             };
+
+            if let Some(grouping_id) = &self.bind_context.aggregate_info.grouping_id_column {
+                if grouping_id.index != column_binding.index {
+                    column_binding.data_type = Box::new(column_binding.data_type.wrap_nullable());
+                }
+            }
+
             return Ok(BoundColumnRef {
                 column: column_binding,
             }

tests/sqllogictests/suites/duckdb/sql/aggregate/group/group_by_grouping_sets.test

@@ -1,3 +1,10 @@
+query I
+select (number % 2) as a from numbers(5) group by grouping sets (a) order by a;
+----
+0
+1
+
+
 statement ok
 drop table if exists t;