feat(functions): add map_insert function (#15567)

* feat(functions): add map_insert function * feat(functions): check ci * feat(functions): domain fix * feat(functions): add new params format in map_insert * feat(functions): cargo fmt * feat(functions): fix NULL mapType in insert * fix * fix tests * fix typos --------- Co-authored-by: baishen <[email protected]>
databendlabs · Nov 5, 2024 · 3797fb6 · 3797fb6
1 parent 24a2c68
commit 3797fb6
Show file tree

Hide file tree

Showing 5 changed files with 275 additions and 0 deletions.
diff --git a/src/query/functions/src/scalars/map.rs b/src/query/functions/src/scalars/map.rs
@@ -34,6 +34,8 @@ use databend_common_expression::types::SimpleDomain;
 use databend_common_expression::types::ValueType;
 use databend_common_expression::vectorize_1_arg;
 use databend_common_expression::vectorize_with_builder_2_arg;
+use databend_common_expression::vectorize_with_builder_3_arg;
+use databend_common_expression::vectorize_with_builder_4_arg;
 use databend_common_expression::ColumnBuilder;
 use databend_common_expression::Function;
 use databend_common_expression::FunctionDomain;
@@ -348,6 +350,92 @@ pub fn register(registry: &mut FunctionRegistry) {
             },
         );
 
+    registry.register_3_arg_core::<NullableType<MapType<GenericType<0>, GenericType<1>>>, GenericType<0>, GenericType<1>, MapType<GenericType<0>, GenericType<1>>, _, _>(
+        "map_insert",
+        |_, map_domain, key_domain, value_domain| {
+            let domain = map_domain
+                .value
+                .as_ref()
+                .map(|box inner_domain| {
+                    inner_domain
+                        .as_ref()
+                        .map(|(inner_key_domain, inner_value_domain)| (inner_key_domain.merge(key_domain), inner_value_domain.merge(value_domain)))
+                        .unwrap_or((key_domain.clone(), value_domain.clone()))
+                });
+            FunctionDomain::Domain(domain)
+        },
+        vectorize_with_builder_3_arg::<NullableType<MapType<GenericType<0>, GenericType<1>>>, GenericType<0>, GenericType<1>, MapType<GenericType<0>, GenericType<1>>>(
+            |map, key, val, output, ctx| {
+                let key_type = &ctx.generics[0];
+                if !check_valid_map_key_type(key_type) {
+                    ctx.set_error(output.len(), format!("map keys can not be {}", key_type));
+                    output.commit_row();
+                    return;
+                }
+                if let Some(map) = map {
+                    for (k, _) in map.iter() {
+                        if k == key {
+                            ctx.set_error(output.len(), format!("map key `{}` duplicate", key));
+                            output.commit_row();
+                            return;
+                        }
+                    }
+                    for (k, v) in map.iter() {
+                        output.put_item((k, v));
+                    }
+                }
+                output.put_item((key, val));
+                output.commit_row();
+        })
+    );
+
+    registry.register_4_arg_core::<NullableType<MapType<GenericType<0>, GenericType<1>>>, GenericType<0>, GenericType<1>, BooleanType, MapType<GenericType<0>, GenericType<1>>, _, _>(
+        "map_insert",
+        |_, map_domain, key_domain, value_domain, _| {
+            let domain = map_domain
+                .value
+                .as_ref()
+                .map(|box inner_domain| {
+                    inner_domain
+                        .as_ref()
+                        .map(|(inner_key_domain, inner_value_domain)| (inner_key_domain.merge(key_domain), inner_value_domain.merge(value_domain)))
+                        .unwrap_or((key_domain.clone(), value_domain.clone()))
+                });
+            FunctionDomain::Domain(domain)
+        },
+        vectorize_with_builder_4_arg::<NullableType<MapType<GenericType<0>, GenericType<1>>>, GenericType<0>, GenericType<1>, BooleanType, MapType<GenericType<0>, GenericType<1>>>(
+            |map, key, val, update_flag, output, ctx| {
+                let key_type = &ctx.generics[0];
+                if !check_valid_map_key_type(key_type) {
+                    ctx.set_error(output.len(), format!("map keys can not be {}", key_type));
+                    output.commit_row();
+                    return;
+                }
+                let mut is_duplicate = false;
+                if let Some(map) = map {
+                    for (k, _) in map.iter() {
+                        if k == key && !update_flag {
+                            ctx.set_error(output.len(), format!("map key `{}` duplicate", key));
+                            output.commit_row();
+                            return;
+                        }
+                    }
+                    for (k, v) in map.iter() {
+                        if k == key {
+                            is_duplicate = true;
+                            output.put_item((k, val.clone()));
+                        } else {
+                            output.put_item((k, v));
+                        }
+                    }
+                }
+                if !is_duplicate {
+                    output.put_item((key, val));
+                }
+                output.commit_row();
+        })
+    );
+
     registry.register_function_factory("map_pick", |_, args_type: &[DataType]| {
         let return_type = check_map_arg_types(args_type)?;
         Some(Arc::new(Function {

diff --git a/src/query/functions/tests/it/scalars/map.rs b/src/query/functions/tests/it/scalars/map.rs
@@ -34,6 +34,7 @@ fn test_map() {
     test_map_delete(file);
     test_map_contains_key(file);
     test_map_pick(file);
+    test_map_insert(file)
 }
 
 fn test_map_cat(file: &mut impl Write) {
@@ -419,3 +420,51 @@ fn test_map_pick(file: &mut impl Write) {
         &columns,
     );
 }
+
+fn test_map_insert(file: &mut impl Write) {
+    run_ast(file, "map_insert({}, 'k1', 'v1')", &[]);
+    run_ast(file, "map_insert({'k1': 'v1'}, 'k2', 'v2')", &[]);
+    run_ast(
+        file,
+        "map_insert({'k1': 'v1', 'k2': 'v2'}, 'k1', 'v10', false)",
+        &[],
+    );
+    run_ast(
+        file,
+        "map_insert({'k1': 'v1', 'k2': 'v2'}, 'k1', 'v10', true)",
+        &[],
+    );
+
+    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_insert(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'k1', 'v10')",
+        &columns,
+    );
+    run_ast(
+        file,
+        "map_insert(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'a', 'v10', true)",
+        &columns,
+    );
+    run_ast(
+        file,
+        "map_insert(map([a_col, b_col, c_col], [d_col, e_col, f_col]), 'a', 'v10', false)",
+        &columns,
+    );
+}
diff --git a/src/query/functions/tests/it/scalars/testdata/function_list.txt b/src/query/functions/tests/it/scalars/testdata/function_list.txt
@@ -2546,6 +2546,8 @@ Functions overloads:
 1 map_contains_key(Map(T0, T1), T0) :: Boolean
 2 map_contains_key(Map(T0, T1) NULL, T0 NULL) :: Boolean NULL
 0 map_delete FACTORY
+0 map_insert(Map(T0, T1) NULL, T0, T1) :: Map(T0, T1)
+1 map_insert(Map(T0, T1) NULL, T0, T1, Boolean) :: Map(T0, T1)
 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