fix(optimizer): keep the rows of a second DISTINCT after WITH DISTINCT on a primary key

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:

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();

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()));
 }
 
@@ -22,10 +30,8 @@ std::shared_ptr<planner::LogicalOperator> 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 +44,8 @@ RemoveUnnecessaryDistinctOptimizer::visitOperatorReplaceSwitch(
     auto& distinct = op->cast<LogicalDistinct>();
     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 +95,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) {

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"
@@ -1047,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";
@@ -1063,6 +1073,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(

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]