feat(query): Support access Redis data from dictionaries via the `dic…

…t_get` function. (#16389) * feat: dict_get to redis. * fix * fix: resolve_dict_get * fix * fix * fix * fix and fmt * feat: add transform_dict_get * feat: add lazy_static * feat: add key's type Number and prepare to update transform more. * update: bind_create_dictionary * update: transform default expr * fix: transform and binder and type_check * update: operator & cache. * fix * fix: operator * fix: transform--operators * update: transform&resolve&argument * fix * update: binder & test. * feat: redis-server & test. * fmt * update: test. * fix: binder & test. * fix: binder & test. * update: errorcode & test * fix * fix * fix * fix * update: binder. * fix * fix: schema & transform * merge * update:test * fix * fix * fix. * update: binder & transform. * fix: dict_get test * update. * update
datafuselabs · Sep 9, 2024 · 9ac1368 · 9ac1368
1 parent a0940db
commit 9ac1368
Show file tree

Hide file tree

Showing 20 changed files with 926 additions and 67 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/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",

diff --git a/src/common/exception/src/exception_code.rs b/src/common/exception/src/exception_code.rs
@@ -386,12 +386,7 @@ build_exceptions! {
     // dictionary
     DictionaryAlreadyExists(3113),
     UnknownDictionary(3114),
-    UnknownDictionaryId(3115),
-    UnsupportedDictionaryOption(3116),
-    UnsupportedDictionarySource(3117),
-    MissingDictionaryOption(3118),
-    WrongDictionaryFieldExpr(3119),
-
+    DictionarySourceError(3115),
     // Procedure
     UnknownProcedure(3130),
     ProcedureAlreadyExists(3131),

diff --git a/src/common/storage/src/lib.rs b/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;
 

diff --git a/src/meta/app/src/schema/dictionary.rs b/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,

diff --git a/src/query/functions/src/lib.rs b/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",

diff --git a/src/query/service/src/pipelines/builders/builder_async_function.rs b/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(())

diff --git a/src/query/service/src/pipelines/processors/transforms/mod.rs b/src/query/service/src/pipelines/processors/transforms/mod.rs
@@ -25,6 +25,7 @@ mod transform_block_compact_no_split;
 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;

diff --git a/src/query/service/src/pipelines/processors/transforms/transform_async_function.rs b/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)
     }
 }