Fix selects with conditions

polypheny · Dec 31, 2024 · 994b645 · 994b645
1 parent c735156
commit 994b645
Showing 1 changed file with 14 additions and 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
@@ -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 ) {