add group_median_table to test group median.

Rachelint · Dec 21, 2024 · c86cb2e · c86cb2e
1 parent 311f82f
commit c86cb2e
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 3 deletions.
diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs
@@ -448,14 +448,14 @@ impl<T: ArrowNumericType + Send> GroupsAccumulator for MedianGroupsAccumulator<T
         // `nulls` for converted `ListArray`
         let nulls = filtered_null_mask(opt_filter, input_array);
 
-        let converted_list_array = Arc::new(ListArray::new(
+        let converted_list_array = ListArray::new(
             Arc::new(Field::new_list_field(self.data_type.clone(), false)),
             offsets,
             Arc::new(values),
             nulls,
-        ));
+        );
 
-        Ok(vec![converted_list_array])
+        Ok(vec![Arc::new(converted_list_array)])
     }
 
     fn supports_convert_to_state(&self) -> bool {

diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt
@@ -67,6 +67,31 @@ statement ok
 CREATE TABLE test (c1 BIGINT,c2 BIGINT) as values
 (0,null), (1,1), (null,1), (3,2), (3,2)
 
+statement ok
+CREATE TABLE group_median_table (
+    col_group STRING,
+    col_i8 TINYINT,
+    col_i16 SMALLINT,
+    col_i32 INT,
+    col_i64 BIGINT,
+    col_u8 TINYINT UNSIGNED,
+    col_u16 SMALLINT UNSIGNED,
+    col_u32 INT UNSIGNED,
+    col_u64 BIGINT UNSIGNED,
+    col_f32 FLOAT,
+    col_f64 DOUBLE,
+    col_f64_nan DOUBLE
+) as VALUES
+( "group0", -128, -32768, -2147483648, arrow_cast(-9223372036854775808,'Int64'), 0,   0,     0,          arrow_cast(0,'UInt64'),   1.1, 1.1, 1.1 ),
+( "group0", -128, -32768, -2147483648, arrow_cast(-9223372036854775808,'Int64'), 0,   0,     0,          arrow_cast(0,'UInt64'),   4.4, 4.4, arrow_cast('NAN','Float64') ),
+( "group0", 100,  100,    100,         arrow_cast(100,'Int64'),                  100, 100,   100,        arrow_cast(100,'UInt64'), 3.3, 3.3, arrow_cast('NAN','Float64') ),
+( "group0", 127,  32767,  2147483647,  arrow_cast(9223372036854775807,'Int64'),  255, 65535, 4294967295, 18446744073709551615,     2.2, 2.2, arrow_cast('NAN','Float64') ),
+( "group1", -128, -32768, -2147483648, arrow_cast(-9223372036854775808,'Int64'), 0,   0,     0,          arrow_cast(0,'UInt64'),   1.1, 1.1, 1.1 ),
+( "group1", -128, -32768, -2147483648, arrow_cast(-9223372036854775808,'Int64'), 0,   0,     0,          arrow_cast(0,'UInt64'),   4.4, 4.4, arrow_cast('NAN','Float64') ),
+( "group1", 100,  100,    100,         arrow_cast(100,'Int64'),                  101, 100,   100,        arrow_cast(100,'UInt64'), 3.3, 3.3, arrow_cast('NAN','Float64') ),
+( "group1", 125,  32766,  2147483646,  arrow_cast(9223372036854775806,'Int64'),  100, 101,   4294967294, arrow_cast(100,'UInt64'), 3.2, 5.5, arrow_cast('NAN','Float64') ),
+( "group1", 127,  32767,  2147483647,  arrow_cast(9223372036854775807,'Int64'),  255, 65535, 4294967295, 18446744073709551615,     2.2, 2.2, arrow_cast('NAN','Float64') )
+
 #######
 # Error tests
 #######