From 994b645cd013b2ff1621f719d5fef7565a5a3be2 Mon Sep 17 00:00:00 2001
From: Tobias Hafner <mail@tobiashafner.ch>
Date: Tue, 31 Dec 2024 10:47:21 +0100
Subject: [PATCH] Fix selects with conditions

---
 .../db/transaction/locking/AlgTreeRewriter.java  | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/core/src/main/java/org/polypheny/db/transaction/locking/AlgTreeRewriter.java b/core/src/main/java/org/polypheny/db/transaction/locking/AlgTreeRewriter.java
index e8447a0916..275776875c 100644
--- a/core/src/main/java/org/polypheny/db/transaction/locking/AlgTreeRewriter.java
+++ b/core/src/main/java/org/polypheny/db/transaction/locking/AlgTreeRewriter.java
@@ -59,6 +59,7 @@
 import org.polypheny.db.algebra.logical.relational.LogicalRelTableFunctionScan;
 import org.polypheny.db.algebra.logical.relational.LogicalRelUnion;
 import org.polypheny.db.algebra.logical.relational.LogicalRelValues;
+import org.polypheny.db.catalog.entity.Entity;
 import org.polypheny.db.transaction.Transaction;
 
 public class AlgTreeRewriter extends AlgShuttleImpl {
@@ -116,13 +117,24 @@ public AlgNode visit( LogicalRelFilter filter ) {
     public AlgNode visit( LogicalRelProject project ) {
         AlgNode input = project.getInput();
         if ( input instanceof LogicalRelScan || input instanceof LogicalRelFilter ) {
-            //TODO TH: some inputs might not have an entity and thus return null here...
-            LogicalRelIdCollector collector = LogicalRelIdCollector.create( input, transaction, input.getEntity() );
+            Entity entity = findEntity( project );
+            LogicalRelIdCollector collector = LogicalRelIdCollector.create( input, transaction, entity );
             return project.copy( project.getTraitSet(), List.of(collector) );
         }
         return visitChildren( project );
     }
 
+    public Entity findEntity(AlgNode node) {
+        Entity entity = null;
+        while (entity == null && node != null) {
+            entity = node.getEntity();
+            if (node.getInputs().isEmpty()) {
+                continue;
+            }
+            node = node.getInput( 0 );
+        }
+        return entity;
+    }
 
     @Override
     public AlgNode visit( LogicalRelJoin join ) {