feat(query): Support access Redis data from dictionaries via the dict_get function.

Cargo.toml

@@ -269,6 +269,7 @@ opendal = { version = "0.49.0", features = [
     "services-moka",
     "services-webhdfs",
     "services-huggingface",
+    "services-redis",
 ] }
 openraft = { git = "https://github.com/drmingdrmer/openraft", tag = "v0.10.0-alpha.6", features = [
     "serde",

scripts/ci/deploy/databend-query-standalone.sh

@@ -31,4 +31,4 @@ echo 'Start databend-query...'
 nohup target/${BUILD_PROFILE}/databend-query -c scripts/ci/deploy/config/databend-query-node-1.toml --internal-enable-sandbox-tenant &
 
 echo "Waiting on databend-query 10 seconds..."
-python3 scripts/ci/wait_tcp.py --timeout 30 --port 8000
+python3 scripts/ci/wait_tcp.py --timeout 30 --port 8000

src/common/exception/src/exception_code.rs

@@ -386,12 +386,6 @@ build_exceptions! {
     // dictionary
     DictionaryAlreadyExists(3113),
     UnknownDictionary(3114),
-    UnknownDictionaryId(3115),
-    UnsupportedDictionaryOption(3116),
-    UnsupportedDictionarySource(3117),
-    MissingDictionaryOption(3118),
-    WrongDictionaryFieldExpr(3119),
-
     // Procedure
     UnknownProcedure(3130),
     ProcedureAlreadyExists(3131),

src/common/storage/src/lib.rs

@@ -36,6 +36,7 @@ pub use config::ShareTableConfig;
 pub use config::StorageConfig;
 
 mod operator;
+pub use operator::build_operator;
 pub use operator::init_operator;
 pub use operator::DataOperator;
 

src/meta/app/src/schema/dictionary.rs

@@ -20,6 +20,8 @@ use std::sync::Arc;
 
 use chrono::DateTime;
 use chrono::Utc;
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
 use databend_common_expression::TableSchema;
 
 use super::dictionary_name_ident::DictionaryNameIdent;
@@ -77,6 +79,47 @@ impl Default for DictionaryMeta {
     }
 }
 
+impl DictionaryMeta {
+    pub fn build_sql_connection_url(&self) -> Result<String> {
+        let username = self
+            .options
+            .get("username")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `username`"))?;
+        let password = self
+            .options
+            .get("password")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `password`"))?;
+        let host = self
+            .options
+            .get("host")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `host`"))?;
+        let port = self
+            .options
+            .get("port")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `port`"))?;
+        let db = self
+            .options
+            .get("db")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `db`"))?;
+        Ok(format!(
+            "mysql://{}:{}@{}:{}/{}",
+            username, password, host, port, db
+        ))
+    }
+
+    pub fn build_redis_connection_url(&self) -> Result<String> {
+        let host = self
+            .options
+            .get("host")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `host`"))?;
+        let port = self
+            .options
+            .get("port")
+            .ok_or_else(|| ErrorCode::BadArguments("Miss option `port`"))?;
+        Ok(format!("tcp://{}:{}", host, port))
+    }
+}
+
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct CreateDictionaryReq {
     pub dictionary_ident: DictionaryNameIdent,

src/query/functions/src/lib.rs

@@ -44,7 +44,7 @@ pub fn is_builtin_function(name: &str) -> bool {
 #[ctor]
 pub static BUILTIN_FUNCTIONS: FunctionRegistry = builtin_functions();
 
-pub const ASYNC_FUNCTIONS: [&str; 1] = ["nextval"];
+pub const ASYNC_FUNCTIONS: [&str; 2] = ["nextval", "dict_get"];
 
 pub const GENERAL_WINDOW_FUNCTIONS: [&str; 13] = [
     "row_number",

src/query/service/src/pipelines/builders/builder_async_function.rs

@@ -23,8 +23,13 @@ impl PipelineBuilder {
     pub(crate) fn build_async_function(&mut self, async_function: &AsyncFunction) -> Result<()> {
         self.build_pipeline(&async_function.input)?;
 
+        let operators = TransformAsyncFunction::init_operators(&async_function.async_func_descs)?;
         self.main_pipeline.add_async_transformer(|| {
-            TransformAsyncFunction::new(self.ctx.clone(), async_function.async_func_descs.clone())
+            TransformAsyncFunction::new(
+                self.ctx.clone(),
+                async_function.async_func_descs.clone(),
+                operators.clone(),
+            )
         });
 
         Ok(())

src/query/service/src/pipelines/processors/transforms/mod.rs

@@ -24,6 +24,7 @@ mod transform_async_function;
 mod transform_cache_scan;
 mod transform_cast_schema;
 mod transform_create_sets;
+mod transform_dictionary;
 mod transform_expression_scan;
 mod transform_filter;
 mod transform_limit;

src/query/service/src/pipelines/processors/transforms/transform_async_function.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use databend_common_exception::Result;
@@ -25,21 +26,29 @@ use databend_common_meta_app::schema::GetSequenceNextValueReq;
 use databend_common_meta_app::schema::SequenceIdent;
 use databend_common_pipeline_transforms::processors::AsyncTransform;
 use databend_common_storages_fuse::TableContext;
+use opendal::Operator;
 
 use crate::sessions::QueryContext;
 use crate::sql::executor::physical_plans::AsyncFunctionDesc;
 use crate::sql::plans::AsyncFunctionArgument;
 
 pub struct TransformAsyncFunction {
     ctx: Arc<QueryContext>,
+    // key is the index of async_func_desc
+    pub(crate) operators: BTreeMap<usize, Arc<Operator>>,
     async_func_descs: Vec<AsyncFunctionDesc>,
 }
 
 impl TransformAsyncFunction {
-    pub fn new(ctx: Arc<QueryContext>, async_func_descs: Vec<AsyncFunctionDesc>) -> Self {
+    pub fn new(
+        ctx: Arc<QueryContext>,
+        async_func_descs: Vec<AsyncFunctionDesc>,
+        operators: BTreeMap<usize, Arc<Operator>>,
+    ) -> Self {
         Self {
             ctx,
             async_func_descs,
+            operators,
         }
     }
 
@@ -80,7 +89,7 @@ impl AsyncTransform for TransformAsyncFunction {
 
     #[async_backtrace::framed]
     async fn transform(&mut self, mut data_block: DataBlock) -> Result<DataBlock> {
-        for async_func_desc in &self.async_func_descs {
+        for (i, async_func_desc) in self.async_func_descs.iter().enumerate() {
             match &async_func_desc.func_arg {
                 AsyncFunctionArgument::SequenceFunction(sequence_name) => {
                     self.transform_sequence(
@@ -90,9 +99,18 @@ impl AsyncTransform for TransformAsyncFunction {
                     )
                     .await?;
                 }
+                AsyncFunctionArgument::DictGetFunction(dict_arg) => {
+                    self.transform_dict_get(
+                        i,
+                        &mut data_block,
+                        dict_arg,
+                        &async_func_desc.arg_indices,
+                        &async_func_desc.data_type,
+                    )
+                    .await?;
+                }
             }
         }
-
         Ok(data_block)
     }
 }

src/query/service/src/pipelines/processors/transforms/transform_dictionary.rs

@@ -0,0 +1,129 @@
+// Copyright 2021 Datafuse Labs
+//
+// 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.
+
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_expression::types::DataType;
+use databend_common_expression::BlockEntry;
+use databend_common_expression::ColumnBuilder;
+use databend_common_expression::DataBlock;
+use databend_common_expression::Scalar;
+use databend_common_expression::ScalarRef;
+use databend_common_expression::Value;
+use databend_common_storage::build_operator;
+use opendal::services::Redis;
+use opendal::Operator;
+
+use crate::pipelines::processors::transforms::TransformAsyncFunction;
+use crate::sql::executor::physical_plans::AsyncFunctionDesc;
+use crate::sql::plans::AsyncFunctionArgument;
+use crate::sql::plans::DictGetFunctionArgument;
+use crate::sql::plans::DictionarySource;
+use crate::sql::IndexType;
+
+impl TransformAsyncFunction {
+    pub fn init_operators(
+        async_func_descs: &[AsyncFunctionDesc],
+    ) -> Result<BTreeMap<usize, Arc<Operator>>> {
+        let mut operators = BTreeMap::new();
+        for (i, async_func_desc) in async_func_descs.iter().enumerate() {
+            if let AsyncFunctionArgument::DictGetFunction(dict_arg) = &async_func_desc.func_arg {
+                match &dict_arg.dict_source {
+                    DictionarySource::Redis(redis_source) => {
+                        let mut builder = Redis::default().endpoint(&redis_source.connection_url);
+                        if let Some(ref username) = redis_source.username {
+                            builder = builder.username(username);
+                        }
+                        if let Some(ref password) = redis_source.password {
+                            builder = builder.password(password);
+                        }
+                        if let Some(db_index) = redis_source.db_index {
+                            builder = builder.db(db_index);
+                        }
+                        let op = build_operator(builder)?;
+                        operators.insert(i, Arc::new(op));
+                    }
+                    DictionarySource::Mysql(_) => {
+                        return Err(ErrorCode::Unimplemented("Mysql source is unsupported"));
+                    }
+                }
+            }
+        }
+        Ok(operators)
+    }
+
+    // transform add dict get column.
+    pub(crate) async fn transform_dict_get(
+        &self,
+        i: usize,
+        data_block: &mut DataBlock,
+        dict_arg: &DictGetFunctionArgument,
+        arg_indices: &[IndexType],
+        data_type: &DataType,
+    ) -> Result<()> {
+        let op = self.operators.get(&i).unwrap().clone();
+
+        // only support one key field.
+        let arg_index = arg_indices[0];
+        let entry = data_block.get_by_offset(arg_index);
+        let value = match &entry.value {
+            Value::Scalar(scalar) => {
+                if let Scalar::String(key) = scalar {
+                    let buffer = op.read(key).await;
+                    match buffer {
+                        Ok(res) => {
+                            let value =
+                                unsafe { String::from_utf8_unchecked(res.current().to_vec()) };
+                            Value::Scalar(Scalar::String(value))
+                        }
+                        Err(_) => Value::Scalar(dict_arg.default_value.clone()),
+                    }
+                } else {
+                    Value::Scalar(dict_arg.default_value.clone())
+                }
+            }
+            Value::Column(column) => {
+                let mut builder = ColumnBuilder::with_capacity(data_type, column.len());
+                for scalar in column.iter() {
+                    if let ScalarRef::String(key) = scalar {
+                        let buffer = op.read(key).await;
+                        match buffer {
+                            Ok(res) => {
+                                let value =
+                                    unsafe { String::from_utf8_unchecked(res.current().to_vec()) };
+                                builder.push(ScalarRef::String(value.as_str()));
+                            }
+                            Err(_) => {
+                                builder.push(dict_arg.default_value.as_ref());
+                            }
+                        };
+                    } else {
+                        builder.push(dict_arg.default_value.as_ref());
+                    }
+                }
+                Value::Column(builder.build())
+            }
+        };
+        let entry = BlockEntry {
+            data_type: data_type.clone(),
+            value,
+        };
+        data_block.add_column(entry);
+
+        Ok(())
+    }
+}