feat(function): Implement map_size function (#15444)

* map_size

* test

Author: forsaken628

src/query/functions/src/scalars/map.rs

@@ -22,6 +22,8 @@ use databend_common_expression::types::GenericType;
 use databend_common_expression::types::MapType;
 use databend_common_expression::types::NullType;
 use databend_common_expression::types::NullableType;
+use databend_common_expression::types::NumberType;
+use databend_common_expression::types::SimpleDomain;
 use databend_common_expression::vectorize_1_arg;
 use databend_common_expression::vectorize_with_builder_2_arg;
 use databend_common_expression::FunctionDomain;
@@ -156,4 +158,16 @@ pub fn register(registry: &mut FunctionRegistry) {
             |map, _| map.values
         ),
     );
+
+    registry.register_1_arg_core::<EmptyMapType, NumberType<u8>, _, _>(
+        "map_size",
+        |_, _| FunctionDomain::Domain(SimpleDomain { min: 0, max: 0 }),
+        |_, _| Value::Scalar(0u8),
+    );
+
+    registry.register_1_arg::<MapType<GenericType<0>, GenericType<1>>, NumberType<u64>, _, _>(
+        "map_size",
+        |_, _| FunctionDomain::Full,
+        |map, _| map.len() as u64,
+    );
 }

src/query/functions/tests/it/scalars/map.rs

@@ -29,6 +29,7 @@ fn test_map() {
     test_get(file);
     test_map_keys(file);
     test_map_values(file);
+    test_map_size(file);
 }
 
 fn test_create(file: &mut impl Write) {
@@ -144,3 +145,32 @@ fn test_map_values(file: &mut impl Write) {
         &columns,
     );
 }
+
+fn test_map_size(file: &mut impl Write) {
+    run_ast(file, "map_size({})", &[]);
+    run_ast(file, "map_size({'a':1,'b':2,'c':3})", &[]);
+    run_ast(file, "map_size({'a':NULL,'b':2,'c':NULL})", &[]);
+
+    let columns = [
+        ("a_col", StringType::from_data(vec!["a", "b", "c"])),
+        ("b_col", StringType::from_data(vec!["d", "e", "f"])),
+        ("c_col", StringType::from_data(vec!["x", "y", "z"])),
+        (
+            "d_col",
+            StringType::from_data_with_validity(vec!["v1", "v2", "v3"], vec![true, true, true]),
+        ),
+        (
+            "e_col",
+            StringType::from_data_with_validity(vec!["v4", "v5", ""], vec![true, true, false]),
+        ),
+        (
+            "f_col",
+            StringType::from_data_with_validity(vec!["v6", "", "v7"], vec![true, false, true]),
+        ),
+    ];
+    run_ast(
+        file,
+        "map_size(map([a_col, b_col, c_col], [d_col, e_col, f_col]))",
+        &columns,
+    );
+}

src/query/functions/tests/it/scalars/testdata/function_list.txt

@@ -2438,6 +2438,9 @@ Functions overloads:
 0 map_keys(Map(Nothing)) :: Array(Nothing)
 1 map_keys(Map(T0, T1)) :: Array(T0)
 2 map_keys(Map(T0, T1) NULL) :: Array(T0) NULL
+0 map_size(Map(Nothing)) :: UInt8
+1 map_size(Map(T0, T1)) :: UInt64
+2 map_size(Map(T0, T1) NULL) :: UInt64 NULL
 0 map_values(Map(Nothing)) :: Array(Nothing)
 1 map_values(Map(T0, T1)) :: Array(T1)
 2 map_values(Map(T0, T1) NULL) :: Array(T1) NULL

src/query/functions/tests/it/scalars/testdata/map.txt

@@ -338,3 +338,57 @@ evaluation (internal):
 +--------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
 
 
+ast            : map_size({})
+raw expr       : map_size(map(array(), array()))
+checked expr   : map_size<Map(Nothing)>(map<Array(Nothing), Array(Nothing)>(array<>(), array<>()))
+optimized expr : 0_u8
+output type    : UInt8
+output domain  : {0..=0}
+output         : 0
+
+
+ast            : map_size({'a':1,'b':2,'c':3})
+raw expr       : map_size(map(array('a', 'b', 'c'), array(1, 2, 3)))
+checked expr   : map_size<T0=String, T1=UInt8><Map(T0, T1)>(map<T0=String, T1=UInt8><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>("a", "b", "c"), array<T0=UInt8><T0, T0, T0>(1_u8, 2_u8, 3_u8)))
+optimized expr : 3_u64
+output type    : UInt64
+output domain  : {3..=3}
+output         : 3
+
+
+ast            : map_size({'a':NULL,'b':2,'c':NULL})
+raw expr       : map_size(map(array('a', 'b', 'c'), array(NULL, 2, NULL)))
+checked expr   : map_size<T0=String, T1=UInt8 NULL><Map(T0, T1)>(map<T0=String, T1=UInt8 NULL><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>("a", "b", "c"), array<T0=UInt8 NULL><T0, T0, T0>(CAST(NULL AS UInt8 NULL), CAST(2_u8 AS UInt8 NULL), CAST(NULL AS UInt8 NULL))))
+optimized expr : 3_u64
+output type    : UInt64
+output domain  : {3..=3}
+output         : 3
+
+
+ast            : map_size(map([a_col, b_col, c_col], [d_col, e_col, f_col]))
+raw expr       : map_size(map(array(a_col::String, b_col::String, c_col::String), array(d_col::String NULL, e_col::String NULL, f_col::String NULL)))
+checked expr   : map_size<T0=String, T1=String NULL><Map(T0, T1)>(map<T0=String, T1=String NULL><Array(T0), Array(T1)>(array<T0=String><T0, T0, T0>(a_col, b_col, c_col), array<T0=String NULL><T0, T0, T0>(d_col, e_col, f_col)))
+evaluation:
++--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
+|        | a_col       | b_col       | c_col       | d_col         | e_col                | f_col                | Output  |
++--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
+| Type   | String      | String      | String      | String NULL   | String NULL          | String NULL          | UInt64  |
+| Domain | {"a"..="c"} | {"d"..="f"} | {"x"..="z"} | {"v1"..="v3"} | {""..="v5"} ∪ {NULL} | {""..="v7"} ∪ {NULL} | Unknown |
+| Row 0  | 'a'         | 'd'         | 'x'         | 'v1'          | 'v4'                 | 'v6'                 | 3       |
+| Row 1  | 'b'         | 'e'         | 'y'         | 'v2'          | 'v5'                 | NULL                 | 3       |
+| Row 2  | 'c'         | 'f'         | 'z'         | 'v3'          | NULL                 | 'v7'                 | 3       |
++--------+-------------+-------------+-------------+---------------+----------------------+----------------------+---------+
+evaluation (internal):
++--------+-----------------------------------------------------------------------------------------------------------------+
+| Column | Data                                                                                                            |
++--------+-----------------------------------------------------------------------------------------------------------------+
+| a_col  | StringColumn { data: 0x616263, offsets: [0, 1, 2, 3] }                                                          |
+| b_col  | StringColumn { data: 0x646566, offsets: [0, 1, 2, 3] }                                                          |
+| c_col  | StringColumn { data: 0x78797a, offsets: [0, 1, 2, 3] }                                                          |
+| d_col  | NullableColumn { column: StringColumn { data: 0x763176327633, offsets: [0, 2, 4, 6] }, validity: [0b_____111] } |
+| e_col  | NullableColumn { column: StringColumn { data: 0x76347635, offsets: [0, 2, 4, 4] }, validity: [0b_____011] }     |
+| f_col  | NullableColumn { column: StringColumn { data: 0x76367637, offsets: [0, 2, 2, 4] }, validity: [0b_____101] }     |
+| Output | UInt64([3, 3, 3])                                                                                               |
++--------+-----------------------------------------------------------------------------------------------------------------+
+
+

tests/sqllogictests/suites/query/functions/02_0074_function_map.test

@@ -12,6 +12,11 @@ select map_keys({}), map_values({})
 ----
 [] []
 
+query TT
+select map_size({})
+----
+0
+
 query TT
 select map_keys({'k1':1,'k2':2,'k3':null}), map_values({'k1':1,'k2':2,'k3':null})
 ----
@@ -30,6 +35,13 @@ select map_keys(col1), map_keys(col2) from t
 ['k5','k6'] ['d','e','f']
 [] NULL
 
+query TT
+select map_size(col1), map_size(col2) from t
+----
+3 2
+2 3
+0 NULL
+
 query TT
 select map_values(col1), map_values(col2) from t
 ----