Updated bloom filter index for map type to work with mapKeys function

kitaisreal · kitaisreal · commit 92fd6b376f16 · 2021-09-22T12:45:48.000+03:00
diff --git a/src/Interpreters/BloomFilter.cpp b/src/Interpreters/BloomFilter.cpp
@@ -1,11 +1,9 @@
 #include <Interpreters/BloomFilter.h>
 #include <city.h>
 #include <Columns/ColumnArray.h>
-#include <Columns/ColumnMap.h>
 #include <Columns/ColumnNullable.h>
 #include <Columns/ColumnLowCardinality.h>
 #include <DataTypes/DataTypeArray.h>
-#include <DataTypes/DataTypeMap.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeLowCardinality.h>
 
@@ -126,14 +124,6 @@ DataTypePtr BloomFilter::getPrimitiveType(const DataTypePtr & data_type)
             throw Exception("Unexpected type " + data_type->getName() + " of bloom filter index.", ErrorCodes::BAD_ARGUMENTS);
     }
 
-    if (const auto * map_type = typeid_cast<const DataTypeMap *>(data_type.get()))
-    {
-        if (!typeid_cast<const DataTypeMap *>(map_type->getKeyType().get()))
-            return getPrimitiveType(map_type->getKeyType());
-        else
-            throw Exception("Unexpected key type " + data_type->getName() + " of bloom filter index for map.", ErrorCodes::BAD_ARGUMENTS);
-    }
-
     if (const auto * nullable_type = typeid_cast<const DataTypeNullable *>(data_type.get()))
         return getPrimitiveType(nullable_type->getNestedType());
 
@@ -148,9 +138,6 @@ ColumnPtr BloomFilter::getPrimitiveColumn(const ColumnPtr & column)
     if (const auto * array_col = typeid_cast<const ColumnArray *>(column.get()))
         return getPrimitiveColumn(array_col->getDataPtr());
 
-    if (const auto * map_col = typeid_cast<const ColumnMap *>(column.get()))
-        return getPrimitiveColumn(map_col->getNestedData().getColumnPtr(0));
-
     if (const auto * nullable_col = typeid_cast<const ColumnNullable *>(column.get()))
         return getPrimitiveColumn(nullable_col->getNestedColumnPtr());
 
diff --git a/src/Interpreters/BloomFilterHash.h b/src/Interpreters/BloomFilterHash.h
@@ -5,7 +5,6 @@
 #include <Columns/IColumn.h>
 #include <Columns/ColumnArray.h>
 #include <Columns/ColumnConst.h>
-#include <Columns/ColumnMap.h>
 #include <Columns/ColumnNullable.h>
 #include <Columns/ColumnsNumber.h>
 #include <Columns/ColumnString.h>
@@ -124,26 +123,6 @@ struct BloomFilterHash
             }
         }
 
-        if (which.isMap())
-        {
-            const auto * map_col = typeid_cast<const ColumnMap *>(column.get());
-            const auto & keys_data = map_col->getNestedData().getColumn(0);
-            if (checkAndGetColumn<ColumnNullable>(keys_data))
-                throw Exception("Unexpected key type " + data_type->getName() + " of bloom filter index.", ErrorCodes::BAD_ARGUMENTS);
-
-            const auto & offsets = map_col->getNestedColumn().getOffsets();
-            limit = offsets[pos + limit - 1] - offsets[pos - 1];  /// PaddedPODArray allows access on index -1.
-            pos = offsets[pos - 1];
-
-            if (limit == 0)
-            {
-                auto index_column = ColumnUInt64::create(1);
-                ColumnUInt64::Container & index_column_vec = index_column->getData();
-                index_column_vec[0] = 0;
-                return index_column;
-            }
-        }
-
         const ColumnPtr actual_col = BloomFilter::getPrimitiveColumn(column);
         const DataTypePtr actual_type = BloomFilter::getPrimitiveType(data_type);
 
diff --git a/src/Storages/MergeTree/MergeTreeIndexConditionBloomFilter.cpp b/src/Storages/MergeTree/MergeTreeIndexConditionBloomFilter.cpp
@@ -512,32 +512,37 @@ bool MergeTreeIndexConditionBloomFilter::traverseASTEquals(
 
         if (function->name == "arrayElement")
         {
+            /** Try to parse arrayElement for mapKeys index.
+              * It is important to ignore keys like column_map['Key'] = '' because if key does not exists in map
+              * we return default value for arrayElement.
+              *
+              * We cannot skip keys that does not exist in map if comparison is with default type value because
+              * that way we skip necessary granules where map key does not exists.
+              */
+            if (value_field == value_type->getDefault())
+                return false;
+
             const auto & col_name = assert_cast<ASTIdentifier *>(function->arguments.get()->children[0].get())->name();
 
-            if (header.has(col_name))
+            auto map_keys_index_column_name = fmt::format("mapKeys({})", col_name);
+
+            if (header.has(map_keys_index_column_name))
             {
-                size_t position = header.getPositionByName(col_name);
+                size_t position = header.getPositionByName(map_keys_index_column_name);
                 const DataTypePtr & index_type = header.getByPosition(position).type;
-                const auto * map_type = typeid_cast<const DataTypeMap *>(index_type.get());
-                if (map_type)
-                {
-                    out.function = function_name == "equals" ? RPNElement::FUNCTION_EQUALS : RPNElement::FUNCTION_NOT_EQUALS;
-
-                    auto & argument = function->arguments.get()->children[1];
-
-                    if (const auto * literal = argument->as<ASTLiteral>())
-                    {
-                        auto element_key = literal->value;
-                        const DataTypePtr actual_type = BloomFilter::getPrimitiveType(index_type);
-                        out.predicate.emplace_back(std::make_pair(position, BloomFilterHash::hashWithField(actual_type.get(), element_key)));
-                    }
-                    else
-                    {
-                        return false;
-                    }
+                out.function = function_name == "equals" ? RPNElement::FUNCTION_EQUALS : RPNElement::FUNCTION_NOT_EQUALS;
 
+                auto & argument = function->arguments.get()->children[1];
+
+                if (const auto * literal = argument->as<ASTLiteral>())
+                {
+                    auto element_key = literal->value;
+                    const DataTypePtr actual_type = BloomFilter::getPrimitiveType(index_type);
+                    out.predicate.emplace_back(std::make_pair(position, BloomFilterHash::hashWithField(actual_type.get(), element_key)));
                     return true;
                 }
+
+                return false;
             }
         }
     }
diff --git a/tests/queries/0_stateless/2021_map_bloom_filter_index.reference b/tests/queries/0_stateless/2021_map_bloom_filter_index.reference
@@ -0,0 +1,12 @@
+Equals with existing key
+0	{'K0':'V0'}
+Equals with non existing key
+Equals with default value
+0	{'K0':'V0'}
+1	{'K1':'V1'}
+Not equals with existing key
+1	{'K1':'V1'}
+Not equals with non existing key
+0	{'K0':'V0'}
+1	{'K1':'V1'}
+Not equals with default value
diff --git a/tests/queries/0_stateless/2021_map_bloom_filter_index.sql b/tests/queries/0_stateless/2021_map_bloom_filter_index.sql
@@ -0,0 +1,30 @@
+DROP TABLE IF EXISTS map_test;
+
+CREATE TABLE map_test
+(
+    row_id UInt32,
+    map Map(String, String),
+    INDEX map_bloom_filter_keys mapKeys(map) TYPE bloom_filter GRANULARITY 1
+) Engine=MergeTree() ORDER BY row_id;
+
+INSERT INTO map_test VALUES (0, {'K0':'V0'}), (1, {'K1':'V1'});
+
+SELECT 'Equals with existing key';
+SELECT * FROM map_test WHERE map['K0'] = 'V0';
+
+SELECT 'Equals with non existing key';
+SELECT * FROM map_test WHERE map['K2'] = 'V2';
+
+SELECT 'Equals with default value';
+SELECT * FROM map_test WHERE map['K3'] = '';
+
+SELECT 'Not equals with existing key';
+SELECT * FROM map_test WHERE map['K0'] != 'V0';
+
+SELECT 'Not equals with non existing key';
+SELECT * FROM map_test WHERE map['K2'] != 'V2';
+
+SELECT 'Not equals with default value';
+SELECT * FROM map_test WHERE map['K3'] != '';
+
+DROP TABLE map_test;
