Merge pull request ClickHouse#81833 from ClickHouse/redis_get_keys

Fix read Redis and KeeperMap by keys

Author: GrigoryPervakov

src/Storages/KVStorageUtils.cpp

@@ -3,9 +3,7 @@
 #include <Columns/ColumnSet.h>
 #include <Columns/ColumnConst.h>
 
-#include <Parsers/ASTFunction.h>
 #include <Parsers/ASTIdentifier.h>
-#include <Parsers/ASTSubquery.h>
 #include <Parsers/ASTLiteral.h>
 #include <Parsers/ASTSelectQuery.h>
 
@@ -15,7 +13,6 @@
 
 #include <Functions/IFunction.h>
 
-
 namespace DB
 {
 
@@ -26,106 +23,8 @@ namespace ErrorCodes
 
 namespace
 {
-// returns keys may be filter by condition
-bool traverseASTFilter(
-    const std::string & primary_key, const DataTypePtr & primary_key_type, const ASTPtr & elem, const PreparedSetsPtr & prepared_sets, const ContextPtr & context, FieldVectorPtr & res)
-{
-    const auto * function = elem->as<ASTFunction>();
-    if (!function)
-        return false;
-
-    if (function->name == "and")
-    {
-        // one child has the key filter condition is ok
-        for (const auto & child : function->arguments->children)
-            if (traverseASTFilter(primary_key, primary_key_type, child, prepared_sets, context, res))
-                return true;
-        return false;
-    }
-    if (function->name == "or")
-    {
-        // make sure every child has the key filter condition
-        for (const auto & child : function->arguments->children)
-            if (!traverseASTFilter(primary_key, primary_key_type, child, prepared_sets, context, res))
-                return false;
-        return true;
-    }
-    if (function->name == "equals" || function->name == "in")
-    {
-        const auto & args = function->arguments->as<ASTExpressionList &>();
-        const ASTIdentifier * ident;
-        std::shared_ptr<IAST> value;
-
-        if (args.children.size() != 2)
-            return false;
-
-        if (function->name == "in")
-        {
-            if (!prepared_sets)
-                return false;
-
-            ident = args.children.at(0)->as<ASTIdentifier>();
-            if (!ident)
-                return false;
-
-            if (ident->name() != primary_key)
-                return false;
-            value = args.children.at(1);
-
-            PreparedSets::Hash set_key = value->getTreeHash(/*ignore_aliases=*/true);
-            FutureSetPtr future_set;
-
-            if ((value->as<ASTSubquery>() || value->as<ASTIdentifier>()))
-                future_set = prepared_sets->findSubquery(set_key);
-            else
-                future_set = prepared_sets->findTuple(set_key, {primary_key_type});
-
-            if (!future_set)
-                return false;
-
-            future_set->buildOrderedSetInplace(context);
-
-            auto set = future_set->get();
-            if (!set)
-                return false;
-
-            if (!set->hasExplicitSetElements())
-                return false;
-
-            set->checkColumnsNumber(1);
-            const auto & set_column = *set->getSetElements()[0];
-
-            if (set_column.getDataType() != primary_key_type->getTypeId())
-                return false;
-
-            for (size_t row = 0; row < set_column.size(); ++row)
-                res->push_back(set_column[row]);
-            return true;
-        }
-
-        if ((ident = args.children.at(0)->as<ASTIdentifier>()))
-            value = args.children.at(1);
-        else if ((ident = args.children.at(1)->as<ASTIdentifier>()))
-            value = args.children.at(0);
-        else
-            return false;
-
-        if (ident->name() != primary_key)
-            return false;
-
-        const auto node = evaluateConstantExpressionAsLiteral(value, context);
-        /// function->name == "equals"
-        if (const auto * literal = node->as<ASTLiteral>())
-        {
-            auto converted_field = convertFieldToType(literal->value, *primary_key_type);
-            if (!converted_field.isNull())
-                res->push_back(converted_field);
-            return true;
-        }
-    }
-    return false;
-}
 
+// returns keys may be filtered by condition
 bool traverseDAGFilter(
     const std::string & primary_key, const DataTypePtr & primary_key_type, const ActionsDAG::Node * elem, const ContextPtr & context, FieldVectorPtr & res)
 {
@@ -239,18 +138,6 @@ std::pair<FieldVectorPtr, bool> getFilterKeys(
     return std::make_pair(res, !matched_keys);
 }
 
-std::pair<FieldVectorPtr, bool> getFilterKeys(
-    const String & primary_key, const DataTypePtr & primary_key_type, const SelectQueryInfo & query_info, const ContextPtr & context)
-{
-    const auto & select = query_info.query->as<ASTSelectQuery &>();
-    if (!select.where())
-        return {{}, true};
-
-    FieldVectorPtr res = std::make_shared<FieldVector>();
-    auto matched_keys = traverseASTFilter(primary_key, primary_key_type, select.where(), query_info.prepared_sets, context, res);
-    return std::make_pair(res, !matched_keys);
-}
-
 std::vector<std::string> serializeKeysToRawString(
     FieldVector::const_iterator & it,
     FieldVector::const_iterator end,

src/Storages/RocksDB/StorageEmbeddedRocksDB.cpp

@@ -6,6 +6,8 @@
 
 #include <Storages/StorageFactory.h>
 #include <Storages/KVStorageUtils.h>
+#include <Storages/AlterCommands.h>
+#include <Storages/RocksDB/RocksDBSettings.h>
 
 #include <Parsers/ASTCreateQuery.h>
 
@@ -29,9 +31,9 @@
 #include <Common/Logger.h>
 #include <Common/logger_useful.h>
 #include <Common/Exception.h>
+#include <Common/JSONBuilder.h>
 #include <Core/Settings.h>
-#include <Storages/AlterCommands.h>
-#include <Storages/RocksDB/RocksDBSettings.h>
+
 #include <IO/SharedThreadPools.h>
 #include <Disks/DiskLocal.h>
 #include <base/sort.h>
@@ -566,6 +568,8 @@ class ReadFromEmbeddedRocksDB : public SourceStepWithFilter
     std::string getName() const override { return "ReadFromEmbeddedRocksDB"; }
     void initializePipeline(QueryPipelineBuilder & pipeline, const BuildQueryPipelineSettings &) override;
     void applyFilters(ActionDAGNodes added_filter_nodes) override;
+    void describeActions(FormatSettings & format_settings) const override;
+    void describeActions(JSONBuilder::JSONMap & map) const override;
 
     ReadFromEmbeddedRocksDB(
         const Names & column_names_,
@@ -666,6 +670,29 @@ void ReadFromEmbeddedRocksDB::applyFilters(ActionDAGNodes added_filter_nodes)
     std::tie(keys, all_scan) = getFilterKeys(storage.primary_key, primary_key_data_type, filter_actions_dag, context);
 }
 
+void ReadFromEmbeddedRocksDB::describeActions(FormatSettings & format_settings) const
+{
+    std::string prefix(format_settings.offset, format_settings.indent_char);
+    if (!all_scan)
+    {
+        format_settings.out << prefix << "ReadType: GetKeys\n";
+        format_settings.out << prefix << "Keys: " << keys->size() << '\n';
+    }
+    else
+        format_settings.out << prefix << "ReadType: FullScan\n";
+}
+
+void ReadFromEmbeddedRocksDB::describeActions(JSONBuilder::JSONMap & map) const
+{
+    if (!all_scan)
+    {
+        map.add("Read Type", "GetKeys");
+        map.add("Keys", keys->size());
+    }
+    else
+        map.add("Read Type", "FullScan");
+}
+
 SinkToStoragePtr StorageEmbeddedRocksDB::write(
     const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, ContextPtr  query_context, bool /*async_insert*/)
 {

src/Storages/SelectQueryInfo.h

@@ -190,6 +190,7 @@ struct SelectQueryInfo
     InputOrderInfoPtr input_order_info;
 
     /// Prepared sets are used for indices by storage engine.
+    /// New analyzer stores prepared sets in planner_context and hashes computed of QueryTree instead of AST.
     /// Example: x IN (1, 2, 3)
     PreparedSetsPtr prepared_sets;