feat: optimize null-eq join fixed keys

windtalker · windtalker · commit 3683cb897b10 · 2026-03-27T14:55:08.000+08:00
diff --git a/dbms/src/Interpreters/Join.cpp b/dbms/src/Interpreters/Join.cpp
@@ -51,17 +51,19 @@ extern const int TYPE_MISMATCH;
 
 namespace
 {
-ColumnRawPtrs getKeyColumns(const Names & key_names, const Block & block)
+ColumnRawPtrs getKeyColumns(const Names & key_names, const Block & block, const std::vector<UInt8> & is_null_eq = {})
 {
     size_t keys_size = key_names.size();
+    RUNTIME_CHECK(is_null_eq.empty() || is_null_eq.size() == keys_size);
     ColumnRawPtrs key_columns(keys_size);
 
     for (size_t i = 0; i < keys_size; ++i)
     {
         key_columns[i] = block.getByName(key_names[i]).column.get();
 
-        /// We will join only keys, where all components are not NULL.
-        if (key_columns[i]->isColumnNullable())
+        /// Ordinary '=' keys join only nested values where all components are not NULL.
+        /// NullEQ keys must keep their nullable wrapper so nullness can participate in key comparison.
+        if (key_columns[i]->isColumnNullable() && (is_null_eq.empty() || is_null_eq[i] == 0))
             key_columns[i] = &static_cast<const ColumnNullable &>(*key_columns[i]).getNestedColumn();
     }
 
@@ -437,12 +439,15 @@ void Join::initBuild(const Block & sample_block, size_t build_concurrency_)
     if (unlikely(initialized))
         throw Exception("Logical error: Join has been initialized", ErrorCodes::LOGICAL_ERROR);
     initialized = true;
-    join_map_method = chooseJoinMapMethod(getKeyColumns(key_names_right, sample_block), key_sizes, collators);
+    join_map_method = chooseJoinMapMethod(
+        getKeyColumns(key_names_right, sample_block, is_null_eq),
+        key_sizes,
+        collators,
+        is_null_eq);
     if (hasNullableNullEqKey(key_names_right, sample_block, is_null_eq))
     {
-        if (join_map_method != JoinMapMethod::serialized)
-            LOG_DEBUG(log, "Force serialized join map method because a nullable NullEQ key is present");
-        join_map_method = JoinMapMethod::serialized;
+        if (join_map_method == JoinMapMethod::serialized)
+            LOG_DEBUG(log, "Use serialized join map method because nullable NullEQ keys do not fit packed fixed keys");
     }
     build_sample_block = sample_block;
     setBuildConcurrencyAndInitJoinPartition(build_concurrency_);
diff --git a/dbms/src/Interpreters/JoinHashMap.cpp b/dbms/src/Interpreters/JoinHashMap.cpp
@@ -14,10 +14,14 @@
 
 #include <Columns/ColumnConst.h>
 #include <Columns/ColumnFixedString.h>
+#include <Columns/ColumnNullable.h>
 #include <Columns/ColumnString.h>
 #include <Common/typeid_cast.h>
+#include <Interpreters/AggregationCommon.h>
 #include <Interpreters/JoinHashMap.h>
 
+#include <limits>
+
 namespace DB
 {
 namespace
@@ -33,31 +37,66 @@ bool canAsColumnString(const IColumn * column)
 JoinMapMethod chooseJoinMapMethod(
     const ColumnRawPtrs & key_columns,
     Sizes & key_sizes,
-    const TiDB::TiDBCollators & collators)
+    const TiDB::TiDBCollators & collators,
+    const std::vector<UInt8> & is_null_eq)
 {
     const size_t keys_size = key_columns.size();
+    RUNTIME_CHECK(is_null_eq.empty() || is_null_eq.size() == keys_size);
 
     if (keys_size == 0)
         return JoinMapMethod::CROSS;
 
+    ColumnRawPtrs nested_key_columns;
+    nested_key_columns.reserve(keys_size);
+    bool has_nullable_null_eq_key = false;
+    for (size_t j = 0; j < keys_size; ++j)
+    {
+        const auto * key_column = key_columns[j];
+        if (const auto * nullable_column = typeid_cast<const ColumnNullable *>(key_column))
+        {
+            nested_key_columns.push_back(&nullable_column->getNestedColumn());
+            has_nullable_null_eq_key = has_nullable_null_eq_key || (!is_null_eq.empty() && is_null_eq[j] != 0);
+        }
+        else
+        {
+            nested_key_columns.push_back(key_column);
+        }
+    }
+
     bool all_fixed = true;
     size_t keys_bytes = 0;
     key_sizes.resize(keys_size);
     for (size_t j = 0; j < keys_size; ++j)
     {
-        if (!key_columns[j]->isFixedAndContiguous())
+        if (!nested_key_columns[j]->isFixedAndContiguous())
         {
             all_fixed = false;
             break;
         }
-        key_sizes[j] = key_columns[j]->sizeOfValueIfFixed();
+        key_sizes[j] = nested_key_columns[j]->sizeOfValueIfFixed();
         keys_bytes += key_sizes[j];
     }
 
+    if (has_nullable_null_eq_key)
+    {
+        if (all_fixed)
+        {
+            if (keys_bytes > (std::numeric_limits<size_t>::max() - std::tuple_size<KeysNullMap<UInt128>>::value))
+                throw Exception("Join: keys sizes overflow", ErrorCodes::LOGICAL_ERROR);
+
+            if (std::tuple_size<KeysNullMap<UInt128>>::value + keys_bytes <= sizeof(UInt128))
+                return JoinMapMethod::nullable_keys128;
+            if (std::tuple_size<KeysNullMap<UInt256>>::value + keys_bytes <= sizeof(UInt256))
+                return JoinMapMethod::nullable_keys256;
+        }
+
+        return JoinMapMethod::serialized;
+    }
+
     /// If there is one numeric key that fits in 64 bits
-    if (keys_size == 1 && key_columns[0]->isNumeric())
+    if (keys_size == 1 && nested_key_columns[0]->isNumeric())
     {
-        size_t size_of_field = key_columns[0]->sizeOfValueIfFixed();
+        size_t size_of_field = nested_key_columns[0]->sizeOfValueIfFixed();
         if (size_of_field == 1)
             return JoinMapMethod::key8;
         if (size_of_field == 2)
@@ -80,7 +119,7 @@ JoinMapMethod chooseJoinMapMethod(
         return JoinMapMethod::keys256;
 
     /// If there is single string key, use hash table of it's values.
-    if (keys_size == 1 && canAsColumnString(key_columns[0]))
+    if (keys_size == 1 && canAsColumnString(nested_key_columns[0]))
     {
         if (collators.empty() || !collators[0])
             return JoinMapMethod::key_strbin;
@@ -108,7 +147,7 @@ JoinMapMethod chooseJoinMapMethod(
         }
     }
 
-    if (keys_size == 1 && typeid_cast<const ColumnFixedString *>(key_columns[0]))
+    if (keys_size == 1 && typeid_cast<const ColumnFixedString *>(nested_key_columns[0]))
         return JoinMapMethod::key_fixed_string;
 
     /// Otherwise, use serialized values as the key.
diff --git a/dbms/src/Interpreters/JoinHashMap.h b/dbms/src/Interpreters/JoinHashMap.h
@@ -143,6 +143,8 @@ struct WithUsedFlag<false, Base> : Base
     M(key_fixed_string)            \
     M(keys128)                     \
     M(keys256)                     \
+    M(nullable_keys128)            \
+    M(nullable_keys256)            \
     M(serialized)
 
 enum class JoinMapMethod
@@ -171,6 +173,8 @@ struct ConcurrentMapsTemplate
     using key_fixed_stringType = ConcurrentHashMapWithSavedHash<StringRef, Mapped>;
     using keys128Type = ConcurrentHashMap<UInt128, Mapped, HashCRC32<UInt128>>;
     using keys256Type = ConcurrentHashMap<UInt256, Mapped, HashCRC32<UInt256>>;
+    using nullable_keys128Type = ConcurrentHashMap<UInt128, Mapped, HashCRC32<UInt128>>;
+    using nullable_keys256Type = ConcurrentHashMap<UInt256, Mapped, HashCRC32<UInt256>>;
     using serializedType = ConcurrentHashMapWithSavedHash<StringRef, Mapped>;
 
     std::unique_ptr<key8Type> key8;
@@ -183,6 +187,8 @@ struct ConcurrentMapsTemplate
     std::unique_ptr<key_fixed_stringType> key_fixed_string;
     std::unique_ptr<keys128Type> keys128;
     std::unique_ptr<keys256Type> keys256;
+    std::unique_ptr<nullable_keys128Type> nullable_keys128;
+    std::unique_ptr<nullable_keys256Type> nullable_keys256;
     std::unique_ptr<serializedType> serialized;
     // TODO: add more cases like Aggregator
 };
@@ -201,6 +207,8 @@ struct MapsTemplate
     using key_fixed_stringType = HashMapWithSavedHash<StringRef, Mapped>;
     using keys128Type = HashMap<UInt128, Mapped, HashCRC32<UInt128>>;
     using keys256Type = HashMap<UInt256, Mapped, HashCRC32<UInt256>>;
+    using nullable_keys128Type = HashMap<UInt128, Mapped, HashCRC32<UInt128>>;
+    using nullable_keys256Type = HashMap<UInt256, Mapped, HashCRC32<UInt256>>;
     using serializedType = HashMapWithSavedHash<StringRef, Mapped>;
 
     std::unique_ptr<key8Type> key8;
@@ -213,6 +221,8 @@ struct MapsTemplate
     std::unique_ptr<key_fixed_stringType> key_fixed_string;
     std::unique_ptr<keys128Type> keys128;
     std::unique_ptr<keys256Type> keys256;
+    std::unique_ptr<nullable_keys128Type> nullable_keys128;
+    std::unique_ptr<nullable_keys256Type> nullable_keys256;
     std::unique_ptr<serializedType> serialized;
     // TODO: add more cases like Aggregator
 };
@@ -230,6 +240,8 @@ struct MapsAny
     using key_fixed_stringType = HashSetWithSavedHash<StringRef>;
     using keys128Type = HashSet<UInt128, HashCRC32<UInt128>>;
     using keys256Type = HashSet<UInt256, HashCRC32<UInt256>>;
+    using nullable_keys128Type = HashSet<UInt128, HashCRC32<UInt128>>;
+    using nullable_keys256Type = HashSet<UInt256, HashCRC32<UInt256>>;
     using serializedType = HashSetWithSavedHash<StringRef>;
 
     std::unique_ptr<key8Type> key8;
@@ -242,6 +254,8 @@ struct MapsAny
     std::unique_ptr<key_fixed_stringType> key_fixed_string;
     std::unique_ptr<keys128Type> keys128;
     std::unique_ptr<keys256Type> keys256;
+    std::unique_ptr<nullable_keys128Type> nullable_keys128;
+    std::unique_ptr<nullable_keys256Type> nullable_keys256;
     std::unique_ptr<serializedType> serialized;
     // TODO: add more cases like Aggregator
 };
@@ -257,5 +271,6 @@ using MapsAllFullWithRowFlag = MapsTemplate<RowRefListWithUsedFlag>; // With fla
 JoinMapMethod chooseJoinMapMethod(
     const ColumnRawPtrs & key_columns,
     Sizes & key_sizes,
-    const TiDB::TiDBCollators & collators);
+    const TiDB::TiDBCollators & collators,
+    const std::vector<UInt8> & is_null_eq = {});
 } // namespace DB
diff --git a/dbms/src/Interpreters/JoinPartition.cpp b/dbms/src/Interpreters/JoinPartition.cpp
@@ -440,6 +440,16 @@ struct KeyGetterForTypeImpl<JoinMapMethod::keys256, Value, Mapped>
     using Type = ColumnsHashing::HashMethodKeysFixed<Value, UInt256, Mapped, false, false>;
 };
 template <typename Value, typename Mapped>
+struct KeyGetterForTypeImpl<JoinMapMethod::nullable_keys128, Value, Mapped>
+{
+    using Type = ColumnsHashing::HashMethodKeysFixed<Value, UInt128, Mapped, true, false>;
+};
+template <typename Value, typename Mapped>
+struct KeyGetterForTypeImpl<JoinMapMethod::nullable_keys256, Value, Mapped>
+{
+    using Type = ColumnsHashing::HashMethodKeysFixed<Value, UInt256, Mapped, true, false>;
+};
+template <typename Value, typename Mapped>
 struct KeyGetterForTypeImpl<JoinMapMethod::serialized, Value, Mapped>
 {
     using Type = ColumnsHashing::HashMethodSerialized<Value, Mapped>;
diff --git a/dbms/src/Interpreters/tests/gtest_join_null_eq.cpp b/dbms/src/Interpreters/tests/gtest_join_null_eq.cpp
@@ -199,9 +199,8 @@ JoinPtr makeSemiJoinTestJoin(ASTTableJoin::Kind kind)
         true);
 }
 
-JoinPtr makeMixedKeyJoin(const std::vector<UInt8> & is_null_eq)
+JoinPtr makeMixedKeyJoin(const std::vector<UInt8> & is_null_eq, const DataTypePtr & key_type)
 {
-    auto nullable_int_type = makeNullable(std::make_shared<DataTypeInt32>());
     auto int_type = std::make_shared<DataTypeInt32>();
     SpillConfig build_spill_config("/tmp", "join_null_eq_build", 0, 0, 0, nullptr);
     SpillConfig probe_spill_config("/tmp", "join_null_eq_probe", 0, 0, 0, nullptr);
@@ -217,11 +216,11 @@ JoinPtr makeMixedKeyJoin(const std::vector<UInt8> & is_null_eq)
         probe_spill_config,
         RestoreConfig{1, 0, 0},
         NamesAndTypes{
-            {mixed_probe_key1_name, nullable_int_type},
-            {mixed_probe_key2_name, nullable_int_type},
+            {mixed_probe_key1_name, key_type},
+            {mixed_probe_key2_name, key_type},
             {mixed_probe_value_name, int_type},
-            {mixed_build_key1_name, nullable_int_type},
-            {mixed_build_key2_name, nullable_int_type},
+            {mixed_build_key1_name, key_type},
+            {mixed_build_key2_name, key_type},
             {mixed_build_value_name, int_type},
         },
         RegisterOperatorSpillContext{},
@@ -236,6 +235,11 @@ JoinPtr makeMixedKeyJoin(const std::vector<UInt8> & is_null_eq)
         true);
 }
 
+JoinPtr makeMixedKeyJoin(const std::vector<UInt8> & is_null_eq)
+{
+    return makeMixedKeyJoin(is_null_eq, makeNullable(std::make_shared<DataTypeInt32>()));
+}
+
 Block makeOuterProbeSampleBlock(bool include_filter = false)
 {
     auto nullable_int_type = makeNullable(std::make_shared<DataTypeInt32>());
@@ -268,28 +272,36 @@ Block makeOuterBuildSampleBlock(bool include_filter = false)
     return block;
 }
 
-Block makeMixedProbeSampleBlock()
+Block makeMixedProbeSampleBlock(const DataTypePtr & key_type)
 {
-    auto nullable_int_type = makeNullable(std::make_shared<DataTypeInt32>());
     auto int_type = std::make_shared<DataTypeInt32>();
     return Block{
-        {nullable_int_type->createColumn(), nullable_int_type, mixed_probe_key1_name},
-        {nullable_int_type->createColumn(), nullable_int_type, mixed_probe_key2_name},
+        {key_type->createColumn(), key_type, mixed_probe_key1_name},
+        {key_type->createColumn(), key_type, mixed_probe_key2_name},
         {int_type->createColumn(), int_type, mixed_probe_value_name},
     };
 }
 
-Block makeMixedBuildSampleBlock()
+Block makeMixedProbeSampleBlock()
+{
+    return makeMixedProbeSampleBlock(makeNullable(std::make_shared<DataTypeInt32>()));
+}
+
+Block makeMixedBuildSampleBlock(const DataTypePtr & key_type)
 {
-    auto nullable_int_type = makeNullable(std::make_shared<DataTypeInt32>());
     auto int_type = std::make_shared<DataTypeInt32>();
     return Block{
-        {nullable_int_type->createColumn(), nullable_int_type, mixed_build_key1_name},
-        {nullable_int_type->createColumn(), nullable_int_type, mixed_build_key2_name},
+        {key_type->createColumn(), key_type, mixed_build_key1_name},
+        {key_type->createColumn(), key_type, mixed_build_key2_name},
         {int_type->createColumn(), int_type, mixed_build_value_name},
     };
 }
 
+Block makeMixedBuildSampleBlock()
+{
+    return makeMixedBuildSampleBlock(makeNullable(std::make_shared<DataTypeInt32>()));
+}
+
 ColumnPtr makeNullableInt32Column(std::initializer_list<std::optional<Int32>> values)
 {
     auto nested = ColumnInt32::create();
@@ -403,13 +415,22 @@ void prepareAndFinalizeMixedJoin(const JoinPtr & join)
 }
 } // namespace
 
-TEST(JoinNullEqTest, NullableNullEqKeyForcesSerializedJoinMapMethod)
+TEST(JoinNullEqTest, NullableNullEqKeyUsesNullablePackedJoinMapMethod)
 {
     auto nullable_int_type = makeNullable(std::make_shared<DataTypeInt32>());
     auto join = makeTestJoin(nullable_int_type, {1});
     join->initBuild(makeSampleBlock(nullable_int_type), 1);
 
-    ASSERT_EQ(join->getJoinMapMethod(), JoinMapMethod::serialized);
+    ASSERT_EQ(join->getJoinMapMethod(), JoinMapMethod::nullable_keys128);
+}
+
+TEST(JoinNullEqTest, NullableMixedNullEqKeysCanUseNullableKeys256JoinMapMethod)
+{
+    auto nullable_int64_type = makeNullable(std::make_shared<DataTypeInt64>());
+    auto join = makeMixedKeyJoin({1, 1}, nullable_int64_type);
+    join->initBuild(makeMixedBuildSampleBlock(nullable_int64_type), 1);
+
+    ASSERT_EQ(join->getJoinMapMethod(), JoinMapMethod::nullable_keys256);
 }
 
 TEST(JoinNullEqTest, DefaultMethodSelectionRemainsForOtherCases)
diff --git a/docs/note/nulleq_join.md b/docs/note/nulleq_join.md
