Functions dictGet, dictHas implicit key cast

src/DataTypes/DataTypesNumber.h

@@ -12,6 +12,7 @@ namespace DB
 template <typename T>
 class DataTypeNumber final : public DataTypeNumberBase<T>
 {
+public:
     bool equals(const IDataType & rhs) const override { return typeid(rhs) == typeid(*this); }
 
     bool canBeUsedAsVersion() const override { return true; }

src/Dictionaries/DictionaryStructure.cpp

@@ -3,6 +3,7 @@
 #include <DataTypes/DataTypeFactory.h>
 #include <DataTypes/DataTypeNullable.h>
 #include <DataTypes/DataTypeArray.h>
+#include <DataTypes/DataTypesNumber.h>
 #include <Functions/FunctionHelpers.h>
 #include <Formats/FormatSettings.h>
 #include <IO/WriteHelpers.h>
@@ -123,25 +124,35 @@ DictionaryStructure::DictionaryStructure(const Poco::Util::AbstractConfiguration
         access_to_key_from_attributes = true;
 }
 
+DataTypes DictionaryStructure::getKeyTypes() const
+{
+    if (id)
+        return {std::make_shared<DataTypeUInt64>()};
+
+    const auto & key_attributes = *key;
+    size_t key_attributes_size = key_attributes.size();
+
+    DataTypes result;
+    result.reserve(key_attributes_size);
+
+    for (size_t i = 0; i < key_attributes_size; ++i)
+        result.emplace_back(key_attributes[i].type);
+
+    return result;
+}
 
 void DictionaryStructure::validateKeyTypes(const DataTypes & key_types) const
 {
+    auto key_attributes_types = getKeyTypes();
+    size_t key_attributes_types_size = key_attributes_types.size();
+
     size_t key_types_size = key_types.size();
-    if (key_types_size != getKeysSize())
+    if (key_types_size != key_attributes_types_size)
         throw Exception(ErrorCodes::TYPE_MISMATCH, "Key structure does not match, expected {}", getKeyDescription());
 
-    if (id && !isUInt64(key_types[0]))
-    {
-        throw Exception(ErrorCodes::TYPE_MISMATCH,
-            "Key type for simple key does not match, expected {}, found {}",
-            std::to_string(0),
-            "UInt64",
-            key_types[0]->getName());
-    }
-
     for (size_t i = 0; i < key_types_size; ++i)
     {
-        const auto & expected_type = (*key)[i].type;
+        const auto & expected_type = key_attributes_types[i];
         const auto & actual_type = key_types[i];
 
         if (!areTypesEqual(expected_type, actual_type))

src/Dictionaries/DictionaryStructure.h

@@ -120,6 +120,7 @@ struct DictionaryStructure final
 
     DictionaryStructure(const Poco::Util::AbstractConfiguration & config, const std::string & config_prefix);
 
+    DataTypes getKeyTypes() const;
     void validateKeyTypes(const DataTypes & key_types) const;
 
     const DictionaryAttribute & getAttribute(const std::string & attribute_name) const;

src/Dictionaries/IDictionary.h

@@ -7,6 +7,7 @@
 #include <Columns/ColumnsNumber.h>
 #include <Interpreters/IExternalLoadable.h>
 #include <Interpreters/StorageID.h>
+#include <Interpreters/castColumn.h>
 #include <Dictionaries/IDictionarySource.h>
 #include <Dictionaries/DictionaryStructure.h>
 #include <DataTypes/IDataType.h>
@@ -17,6 +18,7 @@ namespace DB
 namespace ErrorCodes
 {
     extern const int NOT_IMPLEMENTED;
+    extern const int TYPE_MISMATCH;
 }
 
 class IDictionary;
@@ -107,12 +109,46 @@ class IDictionary : public IExternalLoadable
     virtual bool isInjective(const std::string & attribute_name) const = 0;
 
     /** Subclass must provide key type that is supported by dictionary.
-      * Client will use that key type to provide valid key columns for `getColumn` and `has` functions.
+      * Client will use that key type to provide valid key columns for `getColumn` and `hasKeys` functions.
       */
     virtual DictionaryKeyType getKeyType() const = 0;
 
     virtual DictionarySpecialKeyType getSpecialKeyType() const { return DictionarySpecialKeyType::None; }
 
+    /** Convert key columns for getColumn, hasKeys functions to match dictionary key column types.
+      * Default implementation convert key columns into DictionaryStructure key types.
+      * Subclass can override this method if keys for getColumn, hasKeys functions
+      * are different from DictionaryStructure keys. Or to prevent implicit conversion of key columns.
+      */
+    virtual void convertKeyColumns(Columns & key_columns, DataTypes & key_types) const
+    {
+        const auto & dictionary_structure = getStructure();
+        auto key_attributes_types = dictionary_structure.getKeyTypes();
+        size_t key_attributes_types_size = key_attributes_types.size();
+        size_t key_types_size = key_types.size();
+
+        if (key_types_size != key_attributes_types_size)
+            throw Exception(ErrorCodes::TYPE_MISMATCH,
+                "Dictionary {} key lookup structure does not match, expected {}",
+                getFullName(),
+                dictionary_structure.getKeyDescription());
+
+        for (size_t key_attribute_type_index = 0; key_attribute_type_index < key_attributes_types_size; ++key_attribute_type_index)
+        {
+            const auto & key_attribute_type = key_attributes_types[key_attribute_type_index];
+            auto & key_type = key_types[key_attribute_type_index];
+
+            if (key_attribute_type->equals(*key_type))
+                continue;
+
+            auto & key_column_to_cast = key_columns[key_attribute_type_index];
+            ColumnWithTypeAndName column_to_cast = {key_column_to_cast, key_type, ""};
+            auto casted_column = castColumnAccurate(std::move(column_to_cast), key_attribute_type);
+            key_column_to_cast = std::move(casted_column);
+            key_type = key_attribute_type;
+        }
+    }
+
     /** Subclass must validate key columns and keys types
       * and return column representation of dictionary attribute.
       *

src/Dictionaries/IPAddressDictionary.cpp

@@ -209,6 +209,11 @@ IPAddressDictionary::IPAddressDictionary(
     calculateBytesAllocated();
 }
 
+void IPAddressDictionary::convertKeyColumns(Columns &, DataTypes &) const
+{
+    /// Do not perform any implicit keys conversion for IPAddressDictionary
+}
+
 ColumnPtr IPAddressDictionary::getColumn(
     const std::string & attribute_name,
     const DataTypePtr & result_type,

src/Dictionaries/IPAddressDictionary.h

@@ -69,6 +69,8 @@ class IPAddressDictionary final : public IDictionary
 
     DictionaryKeyType getKeyType() const override { return DictionaryKeyType::Complex; }
 
+    void convertKeyColumns(Columns & key_columns, DataTypes & key_types) const override;
+
     ColumnPtr getColumn(
         const std::string& attribute_name,
         const DataTypePtr & result_type,

src/Dictionaries/PolygonDictionary.cpp

@@ -6,6 +6,7 @@
 #include <Columns/ColumnArray.h>
 #include <Columns/ColumnTuple.h>
 #include <DataTypes/DataTypeArray.h>
+#include <DataTypes/DataTypesNumber.h>
 #include <Functions/FunctionHelpers.h>
 #include <QueryPipeline/Pipe.h>
 #include <Processors/Sources/SourceFromSingleChunk.h>
@@ -41,6 +42,29 @@ IPolygonDictionary::IPolygonDictionary(
     calculateBytesAllocated();
 }
 
+void IPolygonDictionary::convertKeyColumns(Columns & key_columns, DataTypes & key_types) const
+{
+    if (key_columns.size() != 2)
+        throw Exception(ErrorCodes::TYPE_MISMATCH,
+            "Dictionary {} key lookup structure does not match, expected two columns of coordinates with type Float64",
+            getFullName());
+
+    auto float_64_type = std::make_shared<DataTypeFloat64>();
+    size_t key_types_size = key_types.size();
+    for (size_t key_type_index = 0; key_type_index < key_types_size; ++key_type_index)
+    {
+        auto & key_type = key_types[key_type_index];
+        if (float_64_type->equals(*key_type))
+            continue;
+
+        auto & key_column_to_cast = key_columns[key_type_index];
+        ColumnWithTypeAndName column_to_cast = {key_column_to_cast, key_type, ""};
+        auto casted_column = castColumnAccurate(std::move(column_to_cast), float_64_type);
+        key_column_to_cast = std::move(casted_column);
+        key_type = float_64_type;
+    }
+}
+
 ColumnPtr IPolygonDictionary::getColumn(
     const std::string & attribute_name,
     const DataTypePtr & result_type,

src/Dictionaries/PolygonDictionary.h

@@ -95,6 +95,8 @@ class IPolygonDictionary : public IDictionary
 
     DictionaryKeyType getKeyType() const override { return DictionaryKeyType::Complex; }
 
+    void convertKeyColumns(Columns & key_columns, DataTypes & key_types) const override;
+
     ColumnPtr getColumn(
         const std::string& attribute_name,
         const DataTypePtr & result_type,

src/Functions/FunctionsExternalDictionaries.h

@@ -145,6 +145,7 @@ class FunctionDictHas final : public IFunction
     String getName() const override { return name; }
 
     size_t getNumberOfArguments() const override { return 0; }
+
     bool isVariadic() const override { return true; }
 
     bool isDeterministic() const override { return false; }
@@ -183,7 +184,17 @@ class FunctionDictHas final : public IFunction
         if (input_rows_count == 0)
             return result_type->createColumn();
 
-        auto dictionary = helper.getDictionary(arguments[0].column);
+        String dictionary_name;
+
+        if (const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments[0].column.get()))
+            dictionary_name = name_col->getValue<String>();
+        else
+            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
+                "Illegal type {} of first argument of function {}, expected a const string.",
+                arguments[0].type->getName(),
+                getName());
+
+        auto dictionary = helper.getDictionary(dictionary_name);
         auto dictionary_key_type = dictionary->getKeyType();
         auto dictionary_special_key_type = dictionary->getSpecialKeyType();
 
@@ -216,15 +227,8 @@ class FunctionDictHas final : public IFunction
 
         if (dictionary_key_type == DictionaryKeyType::Simple)
         {
-            if (!WhichDataType(key_column_type).isUInt64())
-                 throw Exception(
-                     ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
-                     "Second argument of function {} must be UInt64 when dictionary is simple. Actual type {}.",
-                     getName(),
-                     key_column_with_type.type->getName());
-
             key_columns = {key_column};
-            key_types = {std::make_shared<DataTypeUInt64>()};
+            key_types = {key_column_with_type.type};
         }
         else if (dictionary_key_type == DictionaryKeyType::Complex)
         {
@@ -257,6 +261,8 @@ class FunctionDictHas final : public IFunction
             }
         }
 
+        dictionary->convertKeyColumns(key_columns, key_types);
+
         if (dictionary_special_key_type == DictionarySpecialKeyType::Range)
         {
             key_columns.emplace_back(range_col);
@@ -460,15 +466,8 @@ class FunctionDictGetNoType final : public IFunction
 
         if (dictionary_key_type == DictionaryKeyType::Simple)
         {
-            if (!WhichDataType(key_col_with_type.type).isUInt64())
-                 throw Exception(
-                     ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,
-                     "Third argument of function {} must be UInt64 when dictionary is simple. Actual type {}.",
-                     getName(),
-                     key_col_with_type.type->getName());
-
             key_columns = {key_column};
-            key_types = {std::make_shared<DataTypeUInt64>()};
+            key_types = {key_col_with_type.type};
         }
         else if (dictionary_key_type == DictionaryKeyType::Complex)
         {
@@ -481,7 +480,7 @@ class FunctionDictGetNoType final : public IFunction
                 key_columns = assert_cast<const ColumnTuple &>(*key_column).getColumnsCopy();
                 key_types = assert_cast<const DataTypeTuple &>(*key_column_type).getElements();
             }
-            else if (!isTuple(key_column_type))
+            else
             {
                 size_t keys_size = dictionary->getStructure().getKeysSize();
 
@@ -502,6 +501,8 @@ class FunctionDictGetNoType final : public IFunction
             }
         }
 
+        dictionary->convertKeyColumns(key_columns, key_types);
+
         if (dictionary_special_key_type == DictionarySpecialKeyType::Range)
         {
             key_columns.emplace_back(range_col);

tests/queries/0_stateless/02176_dict_get_has_implicit_key_cast.reference

@@ -0,0 +1,14 @@
+Value
+Value
+Value
+1
+1
+1
+Value
+Value
+Value
+Value
+1
+1
+1
+1

tests/queries/0_stateless/02176_dict_get_has_implicit_key_cast.sql

@@ -0,0 +1,67 @@
+DROP TABLE IF EXISTS 02176_test_simple_key_table;
+CREATE TABLE 02176_test_simple_key_table
+(
+    id UInt64,
+    value String
+) ENGINE=TinyLog;
+
+INSERT INTO 02176_test_simple_key_table VALUES (0, 'Value');
+
+DROP DICTIONARY IF EXISTS 02176_test_simple_key_dictionary;
+CREATE DICTIONARY 02176_test_simple_key_dictionary
+(
+    id UInt64,
+    value String
+)
+PRIMARY KEY id
+SOURCE(CLICKHOUSE(TABLE '02176_test_simple_key_table'))
+LAYOUT(DIRECT());
+
+SELECT dictGet('02176_test_simple_key_dictionary', 'value', toUInt64(0));
+SELECT dictGet('02176_test_simple_key_dictionary', 'value', toUInt8(0));
+SELECT dictGet('02176_test_simple_key_dictionary', 'value', '0');
+SELECT dictGet('02176_test_simple_key_dictionary', 'value', [0]); --{serverError 43}
+
+SELECT dictHas('02176_test_simple_key_dictionary', toUInt64(0));
+SELECT dictHas('02176_test_simple_key_dictionary', toUInt8(0));
+SELECT dictHas('02176_test_simple_key_dictionary', '0');
+SELECT dictHas('02176_test_simple_key_dictionary', [0]); --{serverError 43}
+
+DROP DICTIONARY 02176_test_simple_key_dictionary;
+DROP TABLE 02176_test_simple_key_table;
+
+DROP TABLE IF EXISTS 02176_test_complex_key_table;
+CREATE TABLE 02176_test_complex_key_table
+(
+    id UInt64,
+    id_key String,
+    value String
+) ENGINE=TinyLog;
+
+INSERT INTO 02176_test_complex_key_table VALUES (0, '0', 'Value');
+
+DROP DICTIONARY IF EXISTS 02176_test_complex_key_dictionary;
+CREATE DICTIONARY 02176_test_complex_key_dictionary
+(
+    id UInt64,
+    id_key String,
+    value String
+)
+PRIMARY KEY id, id_key
+SOURCE(CLICKHOUSE(TABLE '02176_test_complex_key_table'))
+LAYOUT(COMPLEX_KEY_DIRECT());
+
+SELECT dictGet('02176_test_complex_key_dictionary', 'value', tuple(toUInt64(0), '0'));
+SELECT dictGet('02176_test_complex_key_dictionary', 'value', tuple(toUInt8(0), '0'));
+SELECT dictGet('02176_test_complex_key_dictionary', 'value', tuple('0', '0'));
+SELECT dictGet('02176_test_complex_key_dictionary', 'value', tuple([0], '0')); --{serverError 43}
+SELECT dictGet('02176_test_complex_key_dictionary', 'value', tuple(toUInt64(0), 0));
+
+SELECT dictHas('02176_test_complex_key_dictionary', tuple(toUInt64(0), '0'));
+SELECT dictHas('02176_test_complex_key_dictionary', tuple(toUInt8(0), '0'));
+SELECT dictHas('02176_test_complex_key_dictionary', tuple('0', '0'));
+SELECT dictHas('02176_test_complex_key_dictionary', tuple([0], '0')); --{serverError 43}
+SELECT dictHas('02176_test_complex_key_dictionary', tuple(toUInt64(0), 0));
+
+DROP DICTIONARY 02176_test_complex_key_dictionary;
+DROP TABLE 02176_test_complex_key_table;