Fixing Not-ready set in constraint.

KochetovNicolai · KochetovNicolai · commit 0174a43d17b6 · 2024-07-09T11:54:21.000Z
diff --git a/src/Processors/Transforms/CheckConstraintsTransform.cpp b/src/Processors/Transforms/CheckConstraintsTransform.cpp
@@ -10,6 +10,7 @@
 #include <Common/FieldVisitorToString.h>
 #include <Common/assert_cast.h>
 #include <Common/quoteString.h>
+#include <Storages/VirtualColumnUtils.h>
 
 
 namespace DB
@@ -31,6 +32,7 @@ CheckConstraintsTransform::CheckConstraintsTransform(
     , table_id(table_id_)
     , constraints_to_check(constraints_.filterConstraints(ConstraintsDescription::ConstraintType::CHECK))
     , expressions(constraints_.getExpressions(context_, header.getNamesAndTypesList()))
+    , context(std::move(context_))
 {
 }
 
@@ -39,6 +41,10 @@ void CheckConstraintsTransform::onConsume(Chunk chunk)
 {
     if (chunk.getNumRows() > 0)
     {
+        if (rows_written == 0)
+            for (const auto & expression : expressions)
+                VirtualColumnUtils::buildSetsForDAG(expression->getActionsDAG(), context);
+
         Block block_to_calculate = getInputPort().getHeader().cloneWithColumns(chunk.getColumns());
         for (size_t i = 0; i < expressions.size(); ++i)
         {
diff --git a/src/Processors/Transforms/CheckConstraintsTransform.h b/src/Processors/Transforms/CheckConstraintsTransform.h
@@ -35,6 +35,7 @@ class CheckConstraintsTransform final : public ExceptionKeepingTransform
     StorageID table_id;
     const ASTs constraints_to_check;
     const ConstraintsExpressions expressions;
+    ContextPtr context;
     size_t rows_written = 0;
     Chunk cur_chunk;
 };
diff --git a/src/Storages/System/StorageSystemDetachedParts.cpp b/src/Storages/System/StorageSystemDetachedParts.cpp
@@ -328,7 +328,7 @@ void ReadFromSystemDetachedParts::applyFilters(ActionDAGNodes added_filter_nodes
 
         filter = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);
         if (filter)
-            VirtualColumnUtils::buildSetsForDAG(filter, context);
+            VirtualColumnUtils::buildSetsForDAG(*filter, context);
     }
 }
 
diff --git a/src/Storages/System/StorageSystemPartsBase.cpp b/src/Storages/System/StorageSystemPartsBase.cpp
@@ -274,7 +274,7 @@ void ReadFromSystemPartsBase::applyFilters(ActionDAGNodes added_filter_nodes)
 
         filter_by_database = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);
         if (filter_by_database)
-            VirtualColumnUtils::buildSetsForDAG(filter_by_database, context);
+            VirtualColumnUtils::buildSetsForDAG(*filter_by_database, context);
 
         block.insert(ColumnWithTypeAndName({}, std::make_shared<DataTypeString>(), table_column_name));
         block.insert(ColumnWithTypeAndName({}, std::make_shared<DataTypeString>(), engine_column_name));
@@ -283,7 +283,7 @@ void ReadFromSystemPartsBase::applyFilters(ActionDAGNodes added_filter_nodes)
 
         filter_by_other_columns = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);
         if (filter_by_other_columns)
-            VirtualColumnUtils::buildSetsForDAG(filter_by_other_columns, context);
+            VirtualColumnUtils::buildSetsForDAG(*filter_by_other_columns, context);
     }
 }
 
diff --git a/src/Storages/VirtualColumnUtils.cpp b/src/Storages/VirtualColumnUtils.cpp
@@ -54,9 +54,9 @@ namespace DB
 namespace VirtualColumnUtils
 {
 
-void buildSetsForDAG(const ActionsDAGPtr & dag, const ContextPtr & context)
+void buildSetsForDAG(const ActionsDAG & dag, const ContextPtr & context)
 {
-    for (const auto & node : dag->getNodes())
+    for (const auto & node : dag.getNodes())
     {
         if (node.type == ActionsDAG::ActionType::COLUMN)
         {
@@ -79,7 +79,7 @@ void buildSetsForDAG(const ActionsDAGPtr & dag, const ContextPtr & context)
 
 void filterBlockWithDAG(ActionsDAGPtr dag, Block & block, ContextPtr context)
 {
-    buildSetsForDAG(dag, context);
+    buildSetsForDAG(*dag, context);
     auto actions = std::make_shared<ExpressionActions>(dag);
     Block block_with_filter = block;
     actions->execute(block_with_filter, /*dry_run=*/ false, /*allow_duplicates_in_input=*/ true);
diff --git a/src/Storages/VirtualColumnUtils.h b/src/Storages/VirtualColumnUtils.h
@@ -26,7 +26,7 @@ void filterBlockWithPredicate(const ActionsDAG::Node * predicate, Block & block,
 void filterBlockWithDAG(ActionsDAGPtr dag, Block & block, ContextPtr context);
 
 /// Builds sets used by ActionsDAG inplace.
-void buildSetsForDAG(const ActionsDAGPtr & dag, const ContextPtr & context);
+void buildSetsForDAG(const ActionsDAG & dag, const ContextPtr & context);
 
 /// Recursively checks if all functions used in DAG are deterministic in scope of query.
 bool isDeterministicInScopeOfQuery(const ActionsDAG::Node * node);
diff --git a/tests/queries/0_stateless/02841_not_ready_set_constraints.reference b/tests/queries/0_stateless/02841_not_ready_set_constraints.reference
@@ -0,0 +1 @@
+42
diff --git a/tests/queries/0_stateless/02841_not_ready_set_constraints.sql b/tests/queries/0_stateless/02841_not_ready_set_constraints.sql
@@ -0,0 +1,19 @@
+DROP TABLE IF EXISTS t1;
+DROP TABLE IF EXISTS t2;
+
+CREATE TABLE t1 (
+    `id` UInt64
+)
+ENGINE = MergeTree ORDER BY id;
+
+INSERT INTO t1(id) VALUES (42);
+
+CREATE TABLE t2 (
+    `conversation` UInt64,
+    CONSTRAINT constraint_conversation CHECK conversation IN (SELECT id FROM t1)
+)
+ENGINE = MergeTree ORDER BY conversation;
+
+INSERT INTO t2(conversation) VALUES (42);
+
+select * from t2;

Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,7 @@`
`10`	`10`	`#include <Common/FieldVisitorToString.h>`
`11`	`11`	`#include <Common/assert_cast.h>`
`12`	`12`	`#include <Common/quoteString.h>`
	`13`	`+#include <Storages/VirtualColumnUtils.h>`
`13`	`14`
`14`	`15`
`15`	`16`	`namespace DB`
`@@ -31,6 +32,7 @@ CheckConstraintsTransform::CheckConstraintsTransform(`
`31`	`32`	`, table_id(table_id_)`
`32`	`33`	`, constraints_to_check(constraints_.filterConstraints(ConstraintsDescription::ConstraintType::CHECK))`
`33`	`34`	`, expressions(constraints_.getExpressions(context_, header.getNamesAndTypesList()))`
	`35`	`+ , context(std::move(context_))`
`34`	`36`	`{`
`35`	`37`	`}`
`36`	`38`
`@@ -39,6 +41,10 @@ void CheckConstraintsTransform::onConsume(Chunk chunk)`
`39`	`41`	`{`
`40`	`42`	`if (chunk.getNumRows() > 0)`
`41`	`43`	`{`
	`44`	`+ if (rows_written == 0)`
	`45`	`+ for (const auto & expression : expressions)`
	`46`	`+ VirtualColumnUtils::buildSetsForDAG(expression->getActionsDAG(), context);`
	`47`	`+`
`42`	`48`	`Block block_to_calculate = getInputPort().getHeader().cloneWithColumns(chunk.getColumns());`
`43`	`49`	`for (size_t i = 0; i < expressions.size(); ++i)`
`44`	`50`	`{`
Original file line number	Diff line number	Diff line change
`@@ -328,7 +328,7 @@ void ReadFromSystemDetachedParts::applyFilters(ActionDAGNodes added_filter_nodes`
`328`	`328`
`329`	`329`	`filter = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);`
`330`	`330`	`if (filter)`
`331`		`- VirtualColumnUtils::buildSetsForDAG(filter, context);`
	`331`	`+ VirtualColumnUtils::buildSetsForDAG(*filter, context);`
`332`	`332`	`}`
`333`	`333`	`}`
`334`	`334`
Original file line number	Diff line number	Diff line change
`@@ -274,7 +274,7 @@ void ReadFromSystemPartsBase::applyFilters(ActionDAGNodes added_filter_nodes)`
`274`	`274`
`275`	`275`	`filter_by_database = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);`
`276`	`276`	`if (filter_by_database)`
`277`		`- VirtualColumnUtils::buildSetsForDAG(filter_by_database, context);`
	`277`	`+ VirtualColumnUtils::buildSetsForDAG(*filter_by_database, context);`
`278`	`278`
`279`	`279`	`block.insert(ColumnWithTypeAndName({}, std::make_shared<DataTypeString>(), table_column_name));`
`280`	`280`	`block.insert(ColumnWithTypeAndName({}, std::make_shared<DataTypeString>(), engine_column_name));`
`@@ -283,7 +283,7 @@ void ReadFromSystemPartsBase::applyFilters(ActionDAGNodes added_filter_nodes)`
`283`	`283`
`284`	`284`	`filter_by_other_columns = VirtualColumnUtils::splitFilterDagForAllowedInputs(predicate, &block);`
`285`	`285`	`if (filter_by_other_columns)`
`286`		`- VirtualColumnUtils::buildSetsForDAG(filter_by_other_columns, context);`
	`286`	`+ VirtualColumnUtils::buildSetsForDAG(*filter_by_other_columns, context);`
`287`	`287`	`}`
`288`	`288`	`}`
`289`	`289`
Original file line number	Diff line number	Diff line change
`@@ -54,9 +54,9 @@ namespace DB`
`54`	`54`	`namespace VirtualColumnUtils`
`55`	`55`	`{`
`56`	`56`
`57`		`-void buildSetsForDAG(const ActionsDAGPtr & dag, const ContextPtr & context)`
	`57`	`+void buildSetsForDAG(const ActionsDAG & dag, const ContextPtr & context)`
`58`	`58`	`{`
`59`		`- for (const auto & node : dag->getNodes())`
	`59`	`+ for (const auto & node : dag.getNodes())`
`60`	`60`	`{`
`61`	`61`	`if (node.type == ActionsDAG::ActionType::COLUMN)`
`62`	`62`	`{`
`@@ -79,7 +79,7 @@ void buildSetsForDAG(const ActionsDAGPtr & dag, const ContextPtr & context)`
`79`	`79`
`80`	`80`	`void filterBlockWithDAG(ActionsDAGPtr dag, Block & block, ContextPtr context)`
`81`	`81`	`{`
`82`		`- buildSetsForDAG(dag, context);`
	`82`	`+ buildSetsForDAG(*dag, context);`
`83`	`83`	`auto actions = std::make_shared<ExpressionActions>(dag);`
`84`	`84`	`Block block_with_filter = block;`
`85`	`85`	`actions->execute(block_with_filter, /dry_run=/ false, /allow_duplicates_in_input=/ true);`