feat: adds regex_extract function and more type tests (#7107)

* feat: adds format, regex_extract function and more type tests

Signed-off-by: Dennis Zhuang <[email protected]>

* fix: forgot functions

Signed-off-by: Dennis Zhuang <[email protected]>

* chore: forgot null type

Signed-off-by: Dennis Zhuang <[email protected]>

* test: forgot date type

Signed-off-by: Dennis Zhuang <[email protected]>

* feat: remove format function

Signed-off-by: Dennis Zhuang <[email protected]>

* test: update results after upgrading datafusion

Signed-off-by: Dennis Zhuang <[email protected]>

---------

Signed-off-by: Dennis Zhuang <[email protected]>

Author: killme2008

Cargo.lock

@@ -191,7 +191,7 @@ prost-types = "0.13"
 raft-engine = { version = "0.4.1", default-features = false }
 rand = "0.9"
 ratelimit = "0.10"
-regex = "1.8"
+regex = "1.12"
 regex-automata = "0.4"
 reqwest = { version = "0.12", default-features = false, features = [
     "json",

Cargo.toml

@@ -36,7 +36,7 @@ object_store_opendal.workspace = true
 orc-rust = { version = "0.6.3", default-features = false, features = ["async"] }
 parquet.workspace = true
 paste.workspace = true
-regex = "1.7"
+regex.workspace = true
 serde.workspace = true
 snafu.workspace = true
 strum.workspace = true

src/common/datasource/Cargo.toml

@@ -51,6 +51,7 @@ nalgebra.workspace = true
 num = "0.4"
 num-traits = "0.2"
 paste.workspace = true
+regex.workspace = true
 s2 = { version = "0.0.12", optional = true }
 serde.workspace = true
 serde_json.workspace = true

src/common/function/Cargo.toml

@@ -34,6 +34,7 @@ use crate::scalars::json::JsonFunction;
 use crate::scalars::matches::MatchesFunction;
 use crate::scalars::matches_term::MatchesTermFunction;
 use crate::scalars::math::MathFunction;
+use crate::scalars::string::register_string_functions;
 use crate::scalars::timestamp::TimestampFunction;
 use crate::scalars::uddsketch_calc::UddSketchCalcFunction;
 use crate::scalars::vector::VectorFunction as VectorScalarFunction;
@@ -154,6 +155,9 @@ pub static FUNCTION_REGISTRY: LazyLock<Arc<FunctionRegistry>> = LazyLock::new(||
     // Json related functions
     JsonFunction::register(&function_registry);
 
+    // String related functions
+    register_string_functions(&function_registry);
+
     // Vector related functions
     VectorScalarFunction::register(&function_registry);
     VectorAggrFunction::register(&function_registry);

src/common/function/src/function_registry.rs

@@ -20,6 +20,7 @@ pub mod json;
 pub mod matches;
 pub mod matches_term;
 pub mod math;
+pub(crate) mod string;
 pub mod vector;
 
 pub(crate) mod hll_count;

src/common/function/src/scalars.rs

@@ -20,7 +20,9 @@ use common_query::error;
 use common_time::{Date, Timestamp};
 use datafusion_common::DataFusionError;
 use datafusion_common::arrow::array::{Array, AsArray, StringViewBuilder};
-use datafusion_common::arrow::datatypes::{ArrowTimestampType, DataType, Date32Type, TimeUnit};
+use datafusion_common::arrow::datatypes::{
+    ArrowTimestampType, DataType, Date32Type, Date64Type, TimeUnit,
+};
 use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, Signature};
 use snafu::ResultExt;
 
@@ -40,6 +42,7 @@ impl Default for DateFormatFunction {
             signature: helper::one_of_sigs2(
                 vec![
                     DataType::Date32,
+                    DataType::Date64,
                     DataType::Timestamp(TimeUnit::Second, None),
                     DataType::Timestamp(TimeUnit::Millisecond, None),
                     DataType::Timestamp(TimeUnit::Microsecond, None),
@@ -115,6 +118,29 @@ impl Function for DateFormatFunction {
                     builder.append_option(result.as_deref());
                 }
             }
+            DataType::Date64 => {
+                let left = left.as_primitive::<Date64Type>();
+                for i in 0..size {
+                    let date = left.is_valid(i).then(|| {
+                        let ms = left.value(i);
+                        Timestamp::new_millisecond(ms)
+                    });
+                    let format = formats.is_valid(i).then(|| formats.value(i));
+
+                    let result = match (date, format) {
+                        (Some(ts), Some(fmt)) => {
+                            Some(ts.as_formatted_string(fmt, Some(timezone)).map_err(|e| {
+                                DataFusionError::Execution(format!(
+                                    "cannot format {ts:?} as '{fmt}': {e}"
+                                ))
+                            })?)
+                        }
+                        _ => None,
+                    };
+
+                    builder.append_option(result.as_deref());
+                }
+            }
             x => {
                 return Err(DataFusionError::Execution(format!(
                     "unsupported input data type {x}"
@@ -137,7 +163,9 @@ mod tests {
     use std::sync::Arc;
 
     use arrow_schema::Field;
-    use datafusion_common::arrow::array::{Date32Array, StringArray, TimestampSecondArray};
+    use datafusion_common::arrow::array::{
+        Date32Array, Date64Array, StringArray, TimestampSecondArray,
+    };
     use datafusion_common::config::ConfigOptions;
     use datafusion_expr::{TypeSignature, Volatility};
 
@@ -166,7 +194,7 @@ mod tests {
                          Signature {
                              type_signature: TypeSignature::OneOf(sigs),
                              volatility: Volatility::Immutable
-                         } if  sigs.len() == 5));
+                         } if  sigs.len() == 6));
     }
 
     #[test]
@@ -213,6 +241,50 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_date64_date_format() {
+        let f = DateFormatFunction::default();
+
+        let dates = vec![Some(123000), None, Some(42000), None];
+        let formats = vec![
+            "%Y-%m-%d %T.%3f",
+            "%Y-%m-%d %T.%3f",
+            "%Y-%m-%d %T.%3f",
+            "%Y-%m-%d %T.%3f",
+        ];
+        let results = [
+            Some("1970-01-01 00:02:03.000"),
+            None,
+            Some("1970-01-01 00:00:42.000"),
+            None,
+        ];
+
+        let mut config_options = ConfigOptions::default();
+        config_options.extensions.insert(FunctionContext::default());
+        let config_options = Arc::new(config_options);
+
+        let args = ScalarFunctionArgs {
+            args: vec![
+                ColumnarValue::Array(Arc::new(Date64Array::from(dates))),
+                ColumnarValue::Array(Arc::new(StringArray::from_iter_values(formats))),
+            ],
+            arg_fields: vec![],
+            number_rows: 4,
+            return_field: Arc::new(Field::new("x", DataType::Utf8View, false)),
+            config_options,
+        };
+        let result = f
+            .invoke_with_args(args)
+            .and_then(|x| x.to_array(4))
+            .unwrap();
+        let vector = result.as_string_view();
+
+        assert_eq!(4, vector.len());
+        for (actual, expect) in vector.iter().zip(results) {
+            assert_eq!(actual, expect);
+        }
+    }
+
     #[test]
     fn test_date_date_format() {
         let f = DateFormatFunction::default();

src/common/function/src/scalars/date/date_format.rs

@@ -0,0 +1,26 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! String scalar functions
+
+mod regexp_extract;
+
+pub(crate) use regexp_extract::RegexpExtractFunction;
+
+use crate::function_registry::FunctionRegistry;
+
+/// Register all string functions
+pub fn register_string_functions(registry: &FunctionRegistry) {
+    RegexpExtractFunction::register(registry);
+}