Repository: ibis-project/ibis-substrait

Issue:
bug: Accessing a `Rel` of a `Join` errors because it gets an expression instead of a relation
### What happened?

I see the following error:
```python
File ".../python3.13/site-packages/ibis_substrait/compiler/translate.py", line 713, in table_column
    field_table = op.rel.values.get(op.name).rel
AttributeError: 'Add' object has no attribute 'rel'
```

I am trying to run the following:
```python
substrait_compiler = SubstraitCompiler()
proto_msg = substrait_compiler.compile(tstat_query.unbind())
```

where `tstat_query` is the following:
```python
(Pdb) tstat_query
r0 := SkyTable: simulated/1_1_1_1_1_500000
  feature_name string
  cell_id      string
  expr_val     int32

r1 := SkyTable: simulated/1_1_1_1_2_500000
  feature_name string
  cell_id      string
  expr_val     int32

r2 := SkyTable: simulated/1_1_1_1_3_500000
  feature_name string
  cell_id      string
  expr_val     int32

r3 := Aggregate[r0]
  groups:
    feature_name: r0.feature_name
  metrics:
    cell_count: CountStar(r0)
    expr_total: Sum(r0.expr_val)
    expr_sumsq: Sum(r0.expr_val * r0.expr_val)

r4 := Aggregate[r1]
  groups:
    feature_name: r1.feature_name
  metrics:
    cell_count: CountStar(r1)
    expr_total: Sum(r1.expr_val)
    expr_sumsq: Sum(r1.expr_val * r1.expr_val)

r5 := Aggregate[r2]
  groups:
    feature_name: r2.feature_name
  metrics:
    cell_count: CountStar(r2)
    expr_total: Sum(r2.expr_val)
    expr_sumsq: Sum(r2.expr_val * r2.expr_val)

r6 := JoinChain[r3]
  JoinLink[inner, r4]
    r3.feature_name == r4.feature_name
  JoinLink[inner, r5]
    r3.feature_name == r5.feature_name
  values:
    feature_name: r3.feature_name
    cell_count:   r3.cell_count + r4.cell_count + r5.cell_count
    expr_total:   r3.expr_total + r4.expr_total + r5.expr_total
    expr_sumsq:   r3.expr_sumsq + r4.expr_sumsq + r5.expr_sumsq

Project[r6]
  feature_name: r6.feature_name
  expr_avg:     r6.expr_total / r6.cell_count
  expr_var:     r6.expr_sumsq / r6.cell_count - r6.expr_total / r6.cell_count * r6.expr_total / r6.cell_count
```

### What version of ibis-substrait are you using?

ibis-substrait==4.0.1
substrait==0.20.0

(note that ibis-substrait 4.0.1 is incompatible with substrait 0.21.0 and 0.22.0 because of change to `JOIN_TYPE_SEMI`)

```python
>>> import ibis_substrait
>>> ibis_substrait.__version__
'4.0.1'

>>> substrait.__version__
'0.20.0'
```python

### What substrait consumer(s) are you using, if any?

this bug is independent of substrait consumer. Though, I will be using duckdb and arrow.

### Relevant log output

```sh
(Pdb) l
708             #
709
710             # List of join reference objects
711             join_tables = op.rel.tables
712             # Join reference containing the field we care about
713  ->         field_table = op.rel.values.get(op.name).rel
714             # Index of that join reference in the list of join references
715             field_table_index = join_tables.index(field_table)
716
717             # Offset by the number of columns in each preceding table
718             join_table_offset = sum(

(Pdb) op
<ibis.expr.operations.relations.Field object at 0x111700ad0>

(Pdb) op.rel
<ibis.expr.operations.relations.JoinChain object at 0x11171bdd0>

(Pdb) op.rel.tables
[<ibis.expr.operations.relations.JoinReference object at 0x11171a630>, <ibis.expr.operations.relations.JoinReference object at 0x11171adb0>, <ibis.expr.operations.relations.JoinReference object at 0x11171bbf0>]

(Pdb) op.name
'expr_total'

(Pdb) op.rel.values
{'feature_name': <ibis.expr.operations.relations.Field object at 0x1117003d0>, 'cell_count': <ibis.expr.operations.numeric.Add object at 0x1115b5150>, 'expr_total': <ibis.expr.operations.numeric.Add object at 0x1115b5d50>, 'expr_sumsq': <ibis.expr.operations.numeric.Add object at 0x1115b7c50>}

(Pdb) op.rel.values.get(op.name)
<ibis.expr.operations.numeric.Add object at 0x1115b5d50>

(Pdb) op.rel.values.get(op.name).values
*** AttributeError: 'Add' object has no attribute 'values'
```

Hints from issue comments:
@drin: I've confirmed that ibis-substrait==4.0.0 works as expected (or at least it produces a result and doesn't throw an error).@tokoko: thanks for raising. #1078 fixed another join-related bug in 4.0.0, but apparently introduced this one as well. I'll take a look.@drin: thanks!

Solution patch:
diff --git a/ibis_substrait/compiler/translate.py b/ibis_substrait/compiler/translate.py
index 4fa3f37..6a9cefa 100644
--- a/ibis_substrait/compiler/translate.py
+++ b/ibis_substrait/compiler/translate.py
@@ -686,58 +686,8 @@ def table_column(
     else:
         base_offset = 0
 
-    if isinstance(op.rel, ops.JoinChain):
-        # JoinChains provide the schema of the joined table (which is great for Ibis)
-        # but for substrait we need the Field index computed with respect to
-        # the original table schemas.  In practice, this means rolling through
-        # the tables in a JoinChain and computing the field index _without_
-        # removing the join key
-        #
-        # Given
-        # Table 1
-        #   a: int
-        #   b: int
-        #
-        # Table 2
-        #   a: int
-        #   c: int
-        #
-        # JoinChain[r0]
-        #  JoinLink[inner, r1]
-        #    r0.a == r1.a
-        #  values:
-        #    a: r0.a
-        #    b: r0.b
-        #    c: r1.c
-        #
-        # If we ask for the field index of `c`, the JoinChain schema will give
-        # us an index of `2`, but it should be `3` because
-        #
-        #  0: table 1 a
-        #  1: table 1 b
-        #  2: table 2 a
-        #  3: table 2 c
-        #
-
-        # List of join reference objects
-        join_tables = op.rel.tables
-        # Join reference containing the field we care about
-        field_table = op.rel.values.get(op.name).rel
-        # Index of that join reference in the list of join references
-        field_table_index = join_tables.index(field_table)
-
-        # Offset by the number of columns in each preceding table
-        join_table_offset = sum(
-            len(join_tables[i].schema) for i in range(field_table_index)
-        )
-        # Then add on the index of the column in the table
-        # Also in the event of renaming due to join collisions, resolve
-        # the renamed column to the original name so we can pull it off the parent table
-        orig_name = op.rel.values[op.name].name
-        relative_offset = join_table_offset + field_table.schema._name_locs[orig_name]
-    else:
-        schema = op.rel.schema
-        relative_offset = schema._name_locs[op.name]
+    schema = op.rel.schema
+    relative_offset = schema._name_locs[op.name]
     absolute_offset = base_offset + relative_offset
     return stalg.Expression(
         selection=stalg.Expression.FieldReference(


Test patch:
diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json
index c8f0adb..64ec236 100644
--- a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json
+++ b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h07/tpc_h07.json
@@ -1055,7 +1055,7 @@
                                                       "selection": {
                                                         "directReference": {
                                                           "structField": {
-                                                            "field": 45
+                                                            "field": 1
                                                           }
                                                         },
                                                         "rootReference": {}
@@ -1087,9 +1087,7 @@
                                                     "value": {
                                                       "selection": {
                                                         "directReference": {
-                                                          "structField": {
-                                                            "field": 41
-                                                          }
+                                                          "structField": {}
                                                         },
                                                         "rootReference": {}
                                                       }
@@ -1135,7 +1133,7 @@
                                                       "selection": {
                                                         "directReference": {
                                                           "structField": {
-                                                            "field": 45
+                                                            "field": 1
                                                           }
                                                         },
                                                         "rootReference": {}
@@ -1167,9 +1165,7 @@
                                                     "value": {
                                                       "selection": {
                                                         "directReference": {
-                                                          "structField": {
-                                                            "field": 41
-                                                          }
+                                                          "structField": {}
                                                         },
                                                         "rootReference": {}
                                                       }
@@ -1209,7 +1205,7 @@
                                       "selection": {
                                         "directReference": {
                                           "structField": {
-                                            "field": 17
+                                            "field": 2
                                           }
                                         },
                                         "rootReference": {}
diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h08/tpc_h08.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h08/tpc_h08.json
index cff5a51..249ab36 100644
--- a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h08/tpc_h08.json
+++ b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h08/tpc_h08.json
@@ -1309,7 +1309,7 @@
                                                           "selection": {
                                                             "directReference": {
                                                               "structField": {
-                                                                "field": 54
+                                                                "field": 3
                                                               }
                                                             },
                                                             "rootReference": {}
@@ -1342,7 +1342,7 @@
                                                           "selection": {
                                                             "directReference": {
                                                               "structField": {
-                                                                "field": 36
+                                                                "field": 4
                                                               }
                                                             },
                                                             "rootReference": {}
@@ -1386,7 +1386,7 @@
                                                   "selection": {
                                                     "directReference": {
                                                       "structField": {
-                                                        "field": 4
+                                                        "field": 5
                                                       }
                                                     },
                                                     "rootReference": {}
diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h09/tpc_h09.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h09/tpc_h09.json
index 2ef4713..9a75b7d 100644
--- a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h09/tpc_h09.json
+++ b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h09/tpc_h09.json
@@ -1119,7 +1119,7 @@
                               "selection": {
                                 "directReference": {
                                   "structField": {
-                                    "field": 29
+                                    "field": 3
                                   }
                                 },
                                 "rootReference": {}
diff --git a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h21/tpc_h21.json b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h21/tpc_h21.json
index ea7540e..d7cd523 100644
--- a/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h21/tpc_h21.json
+++ b/ibis_substrait/tests/compiler/snapshots/test_tpch/test_compile/tpc_h21/tpc_h21.json
@@ -748,7 +748,7 @@
                                                                   "selection": {
                                                                     "directReference": {
                                                                       "structField": {
-                                                                        "field": 25
+                                                                        "field": 1
                                                                       }
                                                                     },
                                                                     "rootReference": {}
@@ -781,7 +781,7 @@
                                                                   "selection": {
                                                                     "directReference": {
                                                                       "structField": {
-                                                                        "field": 19
+                                                                        "field": 2
                                                                       }
                                                                     },
                                                                     "rootReference": {}
@@ -793,7 +793,7 @@
                                                                   "selection": {
                                                                     "directReference": {
                                                                       "structField": {
-                                                                        "field": 18
+                                                                        "field": 3
                                                                       }
                                                                     },
                                                                     "rootReference": {}
@@ -823,7 +823,7 @@
                                                           "selection": {
                                                             "directReference": {
                                                               "structField": {
-                                                                "field": 33
+                                                                "field": 6
                                                               }
                                                             },
                                                             "rootReference": {}
@@ -1053,9 +1053,7 @@
                                                                   "value": {
                                                                     "selection": {
                                                                       "directReference": {
-                                                                        "structField": {
-                                                                          "field": 7
-                                                                        }
+                                                                        "structField": {}
                                                                       },
                                                                       "rootReference": {}
                                                                     }
@@ -1092,7 +1090,7 @@
                                                                     "selection": {
                                                                       "directReference": {
                                                                         "structField": {
-                                                                          "field": 9
+                                                                          "field": 4
                                                                         }
                                                                       },
                                                                       "rootReference": {}
@@ -1344,9 +1342,7 @@
                                                                           "value": {
                                                                             "selection": {
                                                                               "directReference": {
-                                                                                "structField": {
-                                                                                  "field": 7
-                                                                                }
+                                                                                "structField": {}
                                                                               },
                                                                               "rootReference": {}
                                                                             }
@@ -1383,7 +1379,7 @@
                                                                             "selection": {
                                                                               "directReference": {
                                                                                 "structField": {
-                                                                                  "field": 9
+                                                                                  "field": 4
                                                                                 }
                                                                               },
                                                                               "rootReference": {}
diff --git a/ibis_substrait/tests/compiler/test_compiler.py b/ibis_substrait/tests/compiler/test_compiler.py
index 9215458..4c50ee3 100644
--- a/ibis_substrait/tests/compiler/test_compiler.py
+++ b/ibis_substrait/tests/compiler/test_compiler.py
@@ -565,7 +565,7 @@ def test_join_chain_indexing_in_group_by(compiler):
         .root.input.project.input.aggregate.groupings[0]
         .grouping_expressions[0]
         .selection.direct_reference.struct_field.field
-        == 5
+        == 3
     )
 
     expr = join_chain.group_by("c").count().select("c")
@@ -576,7 +576,7 @@ def test_join_chain_indexing_in_group_by(compiler):
         .root.input.project.input.aggregate.groupings[0]
         .grouping_expressions[0]
         .selection.direct_reference.struct_field.field
-        == 3
+        == 2
     )
 
     # Group-by on a column that will be renamed by the joinchain
@@ -588,7 +588,7 @@ def test_join_chain_indexing_in_group_by(compiler):
         .root.input.project.input.aggregate.groupings[0]
         .grouping_expressions[0]
         .selection.direct_reference.struct_field.field
-        == 7
+        == 4
     )