Support NullEQ join keys

dbms/src/Debug/MockExecutor/JoinBinder.cpp

@@ -202,6 +202,7 @@ bool JoinBinder::toTiPBExecutor(
     join->set_join_exec_type(tipb::JoinExecType::TypeHashJoin);
     join->set_inner_idx(inner_index);
     join->set_is_null_aware_semi_join(is_null_aware_semi_join);
+    assert(is_null_eq.empty() || is_null_eq.size() == join_cols.size());
 
     for (const auto & key : join_cols)
     {
@@ -219,6 +220,9 @@ bool JoinBinder::toTiPBExecutor(
             collator_id);
     }
 
+    for (const auto flag : is_null_eq)
+        join->add_is_null_eq(flag != 0);
+
     for (const auto & expr : left_conds)
     {
         tipb::Expr * cond = join->add_left_conditions();
@@ -341,6 +345,7 @@ ExecutorBinderPtr compileJoin(
     ExecutorBinderPtr right,
     tipb::JoinType tp,
     const ASTs & join_cols,
+    const std::vector<UInt8> & is_null_eq,
     const ASTs & left_conds,
     const ASTs & right_conds,
     const ASTs & other_conds,
@@ -359,6 +364,7 @@ ExecutorBinderPtr compileJoin(
         output_schema,
         tp,
         join_cols,
+        is_null_eq,
         left_conds,
         right_conds,
         other_conds,
@@ -407,6 +413,6 @@ ExecutorBinderPtr compileJoin(size_t & executor_index, ExecutorBinderPtr left, E
             join_cols.push_back(key);
         }
     }
-    return compileJoin(executor_index, left, right, tp, join_cols);
+    return compileJoin(executor_index, left, right, tp, join_cols, {});
 }
 } // namespace DB::mock

dbms/src/Debug/MockExecutor/JoinBinder.h

@@ -29,6 +29,7 @@ class JoinBinder : public ExecutorBinder
         const DAGSchema & output_schema_,
         tipb::JoinType tp_,
         const ASTs & join_cols_,
+        const std::vector<UInt8> & is_null_eq_,
         const ASTs & l_conds,
         const ASTs & r_conds,
         const ASTs & o_conds,
@@ -39,6 +40,7 @@ class JoinBinder : public ExecutorBinder
         : ExecutorBinder(index_, "Join_" + std::to_string(index_), output_schema_)
         , tp(tp_)
         , join_cols(join_cols_)
+        , is_null_eq(is_null_eq_)
         , left_conds(l_conds)
         , right_conds(r_conds)
         , other_conds(o_conds)
@@ -77,6 +79,7 @@ class JoinBinder : public ExecutorBinder
     tipb::JoinType tp;
 
     const ASTs join_cols{};
+    const std::vector<UInt8> is_null_eq{};
     const ASTs left_conds{};
     const ASTs right_conds{};
     const ASTs other_conds{};
@@ -93,6 +96,7 @@ ExecutorBinderPtr compileJoin(
     ExecutorBinderPtr right,
     tipb::JoinType tp,
     const ASTs & join_cols,
+    const std::vector<UInt8> & is_null_eq = {},
     const ASTs & left_conds = {},
     const ASTs & right_conds = {},
     const ASTs & other_conds = {},

dbms/src/Flash/Coprocessor/JoinInterpreterHelper.cpp

@@ -188,6 +188,19 @@ JoinKeyTypes getJoinKeyTypes(const tipb::Join & join)
     return join_key_types;
 }
 
+std::vector<UInt8> getJoinKeyNullEqFlags(const tipb::Join & join)
+{
+    if (unlikely(join.is_null_eq_size() != 0 && join.is_null_eq_size() != join.left_join_keys_size()))
+        throw TiFlashException(
+            "size of join.is_null_eq does not match size of join.left_join_keys/right_join_keys",
+            Errors::Coprocessor::BadRequest);
+
+    std::vector<UInt8> is_null_eq(join.left_join_keys_size(), 0);
+    for (int i = 0; i < join.is_null_eq_size(); ++i)
+        is_null_eq[i] = join.is_null_eq(i) ? 1 : 0;
+    return is_null_eq;
+}
+
 TiDB::TiDBCollators getJoinKeyCollators(const tipb::Join & join, const JoinKeyTypes & join_key_types, bool is_test)
 {
     TiDB::TiDBCollators collators;
@@ -217,7 +230,18 @@ TiFlashJoin::TiFlashJoin(const tipb::Join & join_, bool is_test) // NOLINT(cppco
     : join(join_)
     , join_key_types(getJoinKeyTypes(join_))
     , join_key_collators(getJoinKeyCollators(join_, join_key_types, is_test))
+    , is_null_eq(getJoinKeyNullEqFlags(join_))
 {
+    if (unlikely(join.is_null_aware_semi_join()))
+    {
+        for (auto flag : is_null_eq)
+        {
+            if (flag != 0)
+                throw TiFlashException(
+                    "NullEQ join keys are incompatible with null-aware semi join",
+                    Errors::Coprocessor::BadRequest);
+        }
+    }
     std::tie(kind, build_side_index) = getJoinKindAndBuildSideIndex(join);
 }
 
@@ -412,6 +436,37 @@ std::tuple<ExpressionActionsPtr, Names, Names, String> prepareJoin(
     return {chain.getLastActions(), std::move(key_names), std::move(original_key_names), std::move(filter_column_name)};
 }
 
+void simplifyNullEqKeyFlags(
+    std::vector<UInt8> & is_null_eq,
+    const ExpressionActionsPtr & probe_prepare_actions,
+    Names & probe_key_names,
+    const ExpressionActionsPtr & build_prepare_actions,
+    Names & build_key_names)
+{
+    RUNTIME_CHECK(probe_key_names.size() == build_key_names.size());
+    RUNTIME_CHECK(probe_key_names.size() == is_null_eq.size());
+
+    for (size_t i = 0; i < is_null_eq.size(); ++i)
+    {
+        if (is_null_eq[i] == 0)
+            continue;
+
+        const auto & probe_type = probe_prepare_actions->getSampleBlock().getByName(probe_key_names[i]).type;
+        const auto & build_type = build_prepare_actions->getSampleBlock().getByName(build_key_names[i]).type;
+        if (!probe_type->isNullable() || !build_type->isNullable())
+        {
+            is_null_eq[i] = 0;
+            continue;
+        }
+
+        RUNTIME_CHECK_MSG(
+            probe_type->equals(*build_type),
+            "NullEQ key type mismatch after prepareJoin: probe={} build={}",
+            probe_type->getName(),
+            build_type->getName());
+    }
+}
+
 std::vector<RuntimeFilterPtr> TiFlashJoin::genRuntimeFilterList(
     const Context & context,
     const NamesAndTypes & source_columns,
@@ -452,6 +507,16 @@ std::vector<RuntimeFilterPtr> TiFlashJoin::genRuntimeFilterList(
     return result;
 }
 
+bool TiFlashJoin::shouldDisableRuntimeFilter() const
+{
+    for (const auto flag : is_null_eq)
+    {
+        if (flag != 0)
+            return true;
+    }
+    return false;
+}
+
 NamesAndTypes genDAGExpressionAnalyzerSourceColumns(Block block, const NamesAndTypes & tidb_schema)
 {
     /// generate source_columns that is used to compile tipb::Expr, the rule is columns in `tidb_schema`

dbms/src/Flash/Coprocessor/JoinInterpreterHelper.h

@@ -140,6 +140,7 @@ struct TiFlashJoin
 
     JoinKeyTypes join_key_types;
     TiDB::TiDBCollators join_key_collators;
+    std::vector<UInt8> is_null_eq;
 
     /// (cartesian) (anti) left outer semi join.
     bool isLeftOuterSemiFamily() const
@@ -217,6 +218,9 @@ struct TiFlashJoin
         const NamesAndTypes & source_columns,
         const std::unordered_map<String, String> & key_names_map,
         const LoggerPtr & log);
+
+    /// Must be called after simplifyNullEqKeyFlags. A remaining NullEQ key has nullable keys on both sides.
+    bool shouldDisableRuntimeFilter() const;
 };
 
 /// @join_prepare_expr_actions: generates join key columns and join filter column
@@ -230,6 +234,16 @@ std::tuple<ExpressionActionsPtr, Names, Names, String> prepareJoin(
     const JoinKeyTypes & join_key_types,
     const google::protobuf::RepeatedPtrField<tipb::Expr> & filters);
 
+/// A NullEQ key is equivalent to an ordinary equality key when either side cannot be NULL.
+/// Simplify such keys before constructing the join so the regular equality-key path can
+/// strip nullable wrappers and filter NULL rows.
+void simplifyNullEqKeyFlags(
+    std::vector<UInt8> & is_null_eq,
+    const ExpressionActionsPtr & probe_prepare_actions,
+    Names & probe_key_names,
+    const ExpressionActionsPtr & build_prepare_actions,
+    Names & build_key_names);
+
 /// generate source_columns that is used to compile tipb::Expr, the rule is columns in `tidb_schema`
 /// must be the first part of the source_columns
 NamesAndTypes genDAGExpressionAnalyzerSourceColumns(Block block, const NamesAndTypes & tidb_schema);