Fix nullable primary key in final

amosbird · Nov 8, 2023 · b36f9bd · b36f9bd
1 parent 16a7d88
commit b36f9bd
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 17 deletions.
diff --git a/src/Processors/QueryPlan/PartsSplitter.cpp b/src/Processors/QueryPlan/PartsSplitter.cpp
@@ -241,10 +241,6 @@ ASTs buildFilters(const KeyDescription & primary_key, const std::vector<Values>
         ASTs values_ast;
         for (size_t i = 0; i < values.size(); ++i)
         {
-            /// NULL is treated as a terminator for > comparison.
-            if (values[i].isNull())
-                break;
-
             const auto & type = primary_key.data_types.at(i);
 
             // PK may contain functions of the table columns, so we need the actual PK AST with all expressions it contains.
@@ -255,25 +251,31 @@ ASTs buildFilters(const KeyDescription & primary_key, const std::vector<Values>
             if (type->isNullable())
             {
                 pks_ast.push_back(makeASTFunction("isNull", pk_ast));
-                values_ast.push_back(std::make_shared<ASTLiteral>(0));
+                values_ast.push_back(std::make_shared<ASTLiteral>(values[i].isNull() ? 1 : 0));
                 pk_ast = makeASTFunction("assumeNotNull", pk_ast);
             }
 
-            ASTPtr component_ast = std::make_shared<ASTLiteral>(values[i]);
-            auto decayed_type = removeNullable(removeLowCardinality(primary_key.data_types.at(i)));
-            // Values of some types (e.g. Date, DateTime) are stored in columns as numbers and we get them as just numbers from the index.
-            // So we need an explicit Cast for them.
-            if (isColumnedAsNumber(decayed_type->getTypeId()) && !isNumber(decayed_type->getTypeId()))
-                component_ast = makeASTFunction("cast", std::move(component_ast), std::make_shared<ASTLiteral>(decayed_type->getName()));
+            pks_ast.push_back(pk_ast);
 
-            pks_ast.push_back(std::move(pk_ast));
-            values_ast.push_back(std::move(component_ast));
+            // If value is null, the comparison is already complete by looking at the null mask column.
+            // Here we put the pk_ast as a placeholder: (pk_null_mask, pk_ast_not_null) > (value_is_null?, pk_ast_not_null).
+            if (values[i].isNull())
+            {
+                values_ast.push_back(pk_ast);
+            }
+            else
+            {
+                ASTPtr component_ast = std::make_shared<ASTLiteral>(values[i]);
+                auto decayed_type = removeNullable(removeLowCardinality(primary_key.data_types.at(i)));
+                // Values of some types (e.g. Date, DateTime) are stored in columns as numbers and we get them as just numbers from the index.
+                // So we need an explicit Cast for them.
+                if (isColumnedAsNumber(decayed_type->getTypeId()) && !isNumber(decayed_type->getTypeId()))
+                    component_ast = makeASTFunction("cast", std::move(component_ast), std::make_shared<ASTLiteral>(decayed_type->getName()));
+
+                values_ast.push_back(std::move(component_ast));
+            }
         }
 
-        /// It indicates (pk1, ...) > (NULL, ...), which is an always false predicate.
-        if (pks_ast.empty())
-            return std::make_shared<ASTLiteral>(0u);
-
         ASTPtr pk_columns_as_tuple = makeASTFunction("tuple", pks_ast);
         ASTPtr values_as_tuple = makeASTFunction("tuple", values_ast);
 
@@ -350,6 +352,7 @@ Pipes buildPipesForReadingByPKRanges(
         auto & filter_function = filters[i];
         if (!filter_function)
             continue;
+        LOG_DEBUG(&::Poco::Logger::get("amosbird"), "filter_function = {}", serializeAST(*filter_function));
         auto syntax_result = TreeRewriter(context).analyze(filter_function, primary_key.expression->getRequiredColumnsWithTypes());
         auto actions = ExpressionAnalyzer(filter_function, syntax_result, context).getActionsDAG(false);
         reorderColumns(*actions, pipes[i].getHeader(), filter_function->getColumnName());

diff --git a/tests/queries/0_stateless/02867_nullable_primary_key_final.reference b/tests/queries/0_stateless/02867_nullable_primary_key_final.reference
@@ -1,2 +1,3 @@
 2023-09-01	2500000000
 2023-09-01	166167
+10
diff --git a/tests/queries/0_stateless/02867_nullable_primary_key_final.sql b/tests/queries/0_stateless/02867_nullable_primary_key_final.sql
@@ -55,3 +55,11 @@ WHERE f2 = 'x'
 GROUP BY 1;
 
 DROP TABLE t;
+
+CREATE TABLE t (o Nullable(String), p Nullable(String)) ENGINE = ReplacingMergeTree ORDER BY (p, o) SETTINGS allow_nullable_key = 1, index_granularity = 2;
+
+INSERT INTO t SELECT number, NULL FROM numbers(10);
+
+SELECT count() FROM t FINAL;
+
+DROP TABLE t;