fix: handle projections after join chain

ibis_substrait/compiler/translate.py

@@ -815,25 +815,22 @@ def filter(
     )
 
 
-@translate.register(ops.Project)
-def project(
-    op: ops.Project,
-    *,
+def apply_projection(
+    schema_len: int,
+    relation: stalg.Rel,
+    values,
     compiler: SubstraitCompiler,
-    child_rel_field_offsets: Mapping[ops.TableNode, int] | None = None,
-    **kwargs: Any,
-) -> stalg.Rel:
-    relation = translate(
-        op.parent, compiler=compiler, child_rel_field_offsets=child_rel_field_offsets
-    )
-    mapping_counter = itertools.count(len(op.parent.schema))
+    child_rel_field_offsets,
+    kwargs,
+):
+    mapping_counter = itertools.count(schema_len)
 
     return stalg.Rel(
         project=stalg.ProjectRel(
             input=relation,
             common=stalg.RelCommon(
                 emit=stalg.RelCommon.Emit(
-                    output_mapping=[next(mapping_counter) for _ in op.values]
+                    output_mapping=[next(mapping_counter) for _ in values]
                 )
             ),
             expressions=[
@@ -843,12 +840,34 @@ def project(
                     child_rel_field_offsets=child_rel_field_offsets,
                     **kwargs,
                 )
-                for k, v in op.values.items()
+                for k, v in values.items()
             ],
         )
     )
 
 
+@translate.register(ops.Project)
+def project(
+    op: ops.Project,
+    *,
+    compiler: SubstraitCompiler,
+    child_rel_field_offsets: Mapping[ops.TableNode, int] | None = None,
+    **kwargs: Any,
+) -> stalg.Rel:
+    relation = translate(
+        op.parent, compiler=compiler, child_rel_field_offsets=child_rel_field_offsets
+    )
+
+    return apply_projection(
+        schema_len=len(op.parent.schema),
+        relation=relation,
+        values=op.values,
+        compiler=compiler,
+        child_rel_field_offsets=child_rel_field_offsets,
+        kwargs=kwargs,
+    )
+
+
 @translate.register(ops.Sort)
 def sort(
     op: ops.Sort,
@@ -944,7 +963,14 @@ def join(
 
             relation = stalg.Rel(join=rel)
 
-    return relation
+    return apply_projection(
+        schema_len=offset,
+        relation=relation,
+        values=op.values,
+        compiler=compiler,
+        child_rel_field_offsets=child_rel_field_offsets,
+        kwargs=kwargs,
+    )
 
 
 @translate.register(ops.Limit)

ibis_substrait/tests/compiler/parity_utils.py

@@ -62,4 +62,5 @@ def execute(self, plan) -> pa.Table:
             df = df.with_column_renamed(
                 column_name, plan.relations[0].root.names[column_number]
             )
-        return df.to_arrow_table()
+        record_batch = df.collect()
+        return pa.Table.from_batches(record_batch)