From e2c8eadf3860bd5f579e689c0e3ce118bc1b3def Mon Sep 17 00:00:00 2001 From: Darek Date: Sat, 17 Oct 2026 08:52:30 +0200 Subject: [PATCH 1/3] fix(optimizer): remove a primary-key DISTINCT before placing flattens MATCH (n:P) WITH DISTINCT n.id AS v WITH DISTINCT coalesce(v, 0) AS w RETURN w returned one row instead of one per node; CASE ... THEN v and [v] did the same. With runtime checks, these queries and v + 1 failed the !flat assertion in FactorizationGroup::setFlat instead. RemoveUnnecessaryDistinctOptimizer ran after FactorizationRewriter, so a removed DISTINCT left behind the FLATTEN placed for it, and the FLATTEN placed for the DISTINCT above then flattened the same chunk again. The pass and AggKeyDependencyOptimizer, whose dependent-key resolution it needs, now run before FactorizationRewriter, which places flattens only for the operators that remain. Plans that keep their DISTINCTs are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../remove_unnecessary_distinct_optimizer.h | 7 +-- src/optimizer/optimizer.cpp | 11 ++-- .../remove_unnecessary_distinct_optimizer.cpp | 11 ++-- .../remove_unnecessary_distinct.test | 51 +++++++++++++++++++ 4 files changed, 64 insertions(+), 16 deletions(-) create mode 100644 test/test_files/optimizer/remove_unnecessary_distinct.test diff --git a/src/include/optimizer/remove_unnecessary_distinct_optimizer.h b/src/include/optimizer/remove_unnecessary_distinct_optimizer.h index af802f5fe4..f6bf2d5681 100644 --- a/src/include/optimizer/remove_unnecessary_distinct_optimizer.h +++ b/src/include/optimizer/remove_unnecessary_distinct_optimizer.h @@ -19,13 +19,14 @@ namespace optimizer { * single scanned node table (such a key uniquely identifies a row). * 2. The DISTINCT's input has multiplicity 1, i.e. its child subtree is a * linear chain of non-multiplying operators (PROJECTION, FILTER, - * MULTIPLICITY_REDUCER, LIMIT, FLATTEN, ...) ending at a plain + * MULTIPLICITY_REDUCER, LIMIT, ...) ending at a plain * SCAN_NODE_TABLE. Joins, extends, recursive/path scans, UNWIND and * materialized-subplan reads (EXPRESSIONS_SCAN/TABLE_FUNCTION_CALL) can all * duplicate rows, so a DISTINCT above them is kept. * - * This optimizer runs after FactorizationRewriter (the plan is in factorized - * form), so it recomputes the factorized schema after rewrites. + * This optimizer runs before FactorizationRewriter (the plan is flat), so it + * recomputes the flat schema after rewrites, and FactorizationRewriter places + * flattens only for the operators that remain. */ class RemoveUnnecessaryDistinctOptimizer { public: diff --git a/src/optimizer/optimizer.cpp b/src/optimizer/optimizer.cpp index c850c784a0..9f5c9527ac 100644 --- a/src/optimizer/optimizer.cpp +++ b/src/optimizer/optimizer.cpp @@ -363,19 +363,18 @@ void Optimizer::optimize(planner::LogicalPlan* plan, main::ClientContext* contex // Degree top-k rewrites need the LIMIT to be folded into ORDER_BY first. countRelTableOptimizer.rewrite(plan); - auto factorizationRewriter = FactorizationRewriter(); - factorizationRewriter.rewrite(plan); - - // AggKeyDependencyOptimizer doesn't change factorization structure and thus can be put - // after FactorizationRewriter. auto aggKeyDependencyOptimizer = AggKeyDependencyOptimizer(); aggKeyDependencyOptimizer.rewrite(plan); // RemoveUnnecessaryDistinctOptimizer should run after AggKeyDependencyOptimizer - // so that dependent-key resolution has already happened. + // so that dependent-key resolution has already happened, and before + // FactorizationRewriter so that no flatten is placed for a DISTINCT it removes. auto removeUnnecessaryDistinctOptimizer = RemoveUnnecessaryDistinctOptimizer(); removeUnnecessaryDistinctOptimizer.rewrite(plan); + auto factorizationRewriter = FactorizationRewriter(); + factorizationRewriter.rewrite(plan); + // RemoveUnnecessaryOrderByOptimizer removes ORDER BY nodes that do not // affect the result (e.g., before a GROUP BY-less aggregate like COUNT(*)). auto removeUnnecessaryOrderByOptimizer = RemoveUnnecessaryOrderByOptimizer(); diff --git a/src/optimizer/remove_unnecessary_distinct_optimizer.cpp b/src/optimizer/remove_unnecessary_distinct_optimizer.cpp index c95a563f06..2c74b7dcae 100644 --- a/src/optimizer/remove_unnecessary_distinct_optimizer.cpp +++ b/src/optimizer/remove_unnecessary_distinct_optimizer.cpp @@ -22,10 +22,8 @@ std::shared_ptr RemoveUnnecessaryDistinctOptimizer::vi op->setChild(i, visitOperator(op->getChild(i))); } auto result = visitOperatorReplaceSwitch(op); - // This optimizer runs after FactorizationRewriter, so the plan is in - // factorized form. Recompute the factorized schema (computeFlatSchema would - // be incorrect here and would also throw on FLATTEN nodes). - result->computeFactorizedSchema(); + // This optimizer runs before FactorizationRewriter, so the plan is flat. + result->computeFlatSchema(); return result; } @@ -38,8 +36,8 @@ RemoveUnnecessaryDistinctOptimizer::visitOperatorReplaceSwitch( auto& distinct = op->cast(); if (isRedundantDistinct(distinct)) { // Bypass the DISTINCT: return its (already-rewritten) child. The caller - // recomputes the factorized schema for the parent, so the parent picks - // up the child's schema cleanly. + // recomputes the flat schema for the parent, so the parent picks up the + // child's schema cleanly. return op->getChild(0); } return op; @@ -89,7 +87,6 @@ bool RemoveUnnecessaryDistinctOptimizer::isChildMultiplicityOne( case LogicalOperatorType::FILTER: case LogicalOperatorType::MULTIPLICITY_REDUCER: case LogicalOperatorType::LIMIT: - case LogicalOperatorType::FLATTEN: case LogicalOperatorType::NODE_LABEL_FILTER: case LogicalOperatorType::NOOP: if (cur->getNumChildren() != 1) { diff --git a/test/test_files/optimizer/remove_unnecessary_distinct.test b/test/test_files/optimizer/remove_unnecessary_distinct.test new file mode 100644 index 0000000000..bfd71dab48 --- /dev/null +++ b/test/test_files/optimizer/remove_unnecessary_distinct.test @@ -0,0 +1,51 @@ +-DATASET CSV empty + +-- + +-CASE ExpressionOverRemovedDistinct +-STATEMENT CREATE NODE TABLE P(id INT64 PRIMARY KEY, flag BOOLEAN) +---- ok +-STATEMENT CREATE (:P {id: 1, flag: true}), (:P {id: 2, flag: false}), (:P {id: 3, flag: true}), (:P {id: 4}) +---- ok +-LOG n.flag depends on the key n.id, so the first DISTINCT has only primary keys and is removed +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v, n.flag AS f WITH DISTINCT CASE WHEN f THEN v ELSE 1 END AS w RETURN w +---- 2 +1 +3 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v WITH DISTINCT coalesce(v, 0) AS w RETURN w +---- 4 +1 +2 +3 +4 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v RETURN DISTINCT coalesce(v, 0) AS w +---- 4 +1 +2 +3 +4 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v WITH DISTINCT [v] AS w RETURN w +---- 4 +[1] +[2] +[3] +[4] +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v WITH DISTINCT v + 1 AS w RETURN w +---- 4 +2 +3 +4 +5 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v MATCH (m:P) WHERE m.id = v + 1 RETURN count(*) +---- 1 +3 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v UNWIND [v, v + 1] AS x RETURN count(DISTINCT x) +---- 1 +5 +-STATEMENT MATCH (n:P) WITH DISTINCT n.id AS v RETURN [v] AS w UNION ALL MATCH (n:P) WHERE n.id = 1 RETURN [n.id] AS w +---- 5 +[1] +[1] +[2] +[3] +[4] From 6efd2e0fb9364defb09d43c4c694112a2075cce5 Mon Sep 17 00:00:00 2001 From: Darek Date: Fri, 23 Oct 2026 23:40:02 +0200 Subject: [PATCH 2/3] chore(optimizer): assert that the DISTINCT removal pass gets a flat plan LogicalFlatten::computeFlatSchema throws, so RemoveUnnecessaryDistinctOptimizer has to run before FactorizationRewriter. With runtime checks the pass now asserts that no FLATTEN is in the plan when it starts, so a later reorder fails at the pass entry instead of inside a FLATTEN's schema code. Release builds skip the walk. Co-Authored-By: Claude Opus 5.5 --- .../remove_unnecessary_distinct_optimizer.cpp | 8 ++++++++ test/optimizer/optimizer_test.cpp | 20 +++++++++++++++++++ 2 files changed, 28 insertions(+) diff --git a/src/optimizer/remove_unnecessary_distinct_optimizer.cpp b/src/optimizer/remove_unnecessary_distinct_optimizer.cpp index 2c74b7dcae..431d6043ad 100644 --- a/src/optimizer/remove_unnecessary_distinct_optimizer.cpp +++ b/src/optimizer/remove_unnecessary_distinct_optimizer.cpp @@ -1,6 +1,8 @@ #include "optimizer/remove_unnecessary_distinct_optimizer.h" #include "binder/expression/property_expression.h" +#include "common/assert.h" +#include "optimizer/logical_operator_collector.h" #include "planner/operator/logical_distinct.h" #include "planner/operator/logical_operator.h" @@ -12,6 +14,12 @@ namespace lbug { namespace optimizer { void RemoveUnnecessaryDistinctOptimizer::rewrite(LogicalPlan* plan) { + // LogicalFlatten::computeFlatSchema throws, so this pass must run before FactorizationRewriter. + RUNTIME_CHECK({ + auto flattens = LogicalFlattenCollector(); + flattens.collect(plan->getLastOperator().get()); + DASSERT(!flattens.hasOperators()); + }); plan->setLastOperator(visitOperator(plan->getLastOperator())); } diff --git a/test/optimizer/optimizer_test.cpp b/test/optimizer/optimizer_test.cpp index ec92fbba76..305da573b4 100644 --- a/test/optimizer/optimizer_test.cpp +++ b/test/optimizer/optimizer_test.cpp @@ -11,8 +11,10 @@ #include "catalog/catalog_entry/node_table_catalog_entry.h" #include "catalog/catalog_entry/rel_group_catalog_entry.h" #include "common/enums/storage_format.h" +#include "common/exception/internal.h" #include "common/system_config.h" #include "graph_test/private_graph_test.h" +#include "optimizer/remove_unnecessary_distinct_optimizer.h" #include "planner/join_order/cost_model.h" #include "planner/operator/logical_filter.h" #include "planner/operator/logical_plan_util.h" @@ -1063,6 +1065,24 @@ TEST_F(OptimizerTest, RemoveUnnecessaryDistinctOnPrimaryKey) { ASSERT_TRUE(plan3 != nullptr); } +TEST_F(OptimizerTest, RemoveUnnecessaryDistinctRequiresFlatPlan) { + // getRoot returns the optimized plan, so FactorizationRewriter has placed its flattens. + auto plan = getRoot("MATCH (a:person)-[:knows]->(b:person) RETURN DISTINCT a.fName, b.fName"); + ASSERT_TRUE( + hasOperatorType(plan->getLastOperator().get(), planner::LogicalOperatorType::FLATTEN)); + try { + optimizer::RemoveUnnecessaryDistinctOptimizer().rewrite(plan.get()); + FAIL() << "the pass accepted a factorized plan"; + } catch (const common::InternalException& e) { +#if defined(RUNTIME_CHECKS) || !defined(NDEBUG) + // The entry check reports it, not LogicalFlatten::computeFlatSchema further down. + EXPECT_NE(std::string(e.what()).find("remove_unnecessary_distinct_optimizer.cpp"), + std::string::npos) + << e.what(); +#endif + } +} + TEST_F(OptimizerTest, CountPrimaryKeyReadsInternalID) { auto scansKey = [&](const std::string& aggregate) { auto plan = getRoot( From 3531dd21caa1f0fc75c84fec3bd6e72e8457e883 Mon Sep 17 00:00:00 2001 From: Darek Date: Sat, 24 Oct 2026 00:32:54 +0200 Subject: [PATCH 3/3] test(optimizer): pin that a removed primary-key DISTINCT leaves no FLATTEN The e2e test catches wrong rows, but a FLATTEN left behind by a removed DISTINCT can also keep the rows right and only make the scan run row by row. Assert that MATCH (n:person) RETURN DISTINCT n.ID and a count(*) over WITH DISTINCT n.ID plan without DISTINCT and without FLATTEN. Co-Authored-By: Claude Opus 5.5 --- test/optimizer/optimizer_test.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test/optimizer/optimizer_test.cpp b/test/optimizer/optimizer_test.cpp index 305da573b4..2982760e17 100644 --- a/test/optimizer/optimizer_test.cpp +++ b/test/optimizer/optimizer_test.cpp @@ -1049,6 +1049,14 @@ TEST_F(OptimizerTest, RemoveUnnecessaryDistinctOnPrimaryKey) { ASSERT_FALSE( hasOperatorType(plan1->getLastOperator().get(), planner::LogicalOperatorType::DISTINCT)) << "DISTINCT on primary key should be removed"; + // A removed DISTINCT leaves no FLATTEN, so the scan stays vectorized. + ASSERT_FALSE( + hasOperatorType(plan1->getLastOperator().get(), planner::LogicalOperatorType::FLATTEN)); + auto plan4 = getRoot("MATCH (n:person) WITH DISTINCT n.ID AS id RETURN count(*)"); + ASSERT_FALSE( + hasOperatorType(plan4->getLastOperator().get(), planner::LogicalOperatorType::DISTINCT)); + ASSERT_FALSE( + hasOperatorType(plan4->getLastOperator().get(), planner::LogicalOperatorType::FLATTEN)); // Sanity check: DISTINCT on a non-primary-key column should be kept auto q2 = "MATCH (n:person) RETURN DISTINCT n.age";