perf(struct): drop unused fields on Struct.__getitem__ (#10996)

Co-authored-by: Phillip Cloud <[email protected]>

Author: NickCrews

ibis/backends/tests/sql/snapshots/test_sql/test_struct_field_simplified/getattr/out.sql

@@ -0,0 +1,2 @@
+SELECT
+  123 + 1 + 1 + 1 + 1 + 1 AS "Add(Add(Add(Add(Add(123, 1), 1), 1), 1), 1)"

ibis/backends/tests/sql/snapshots/test_sql/test_struct_field_simplified/getitem/out.sql

@@ -0,0 +1,2 @@
+SELECT
+  123 + 1 + 1 + 1 + 1 + 1 AS "Add(Add(Add(Add(Add(123, 1), 1), 1), 1), 1)"

ibis/backends/tests/sql/test_sql.py

@@ -663,3 +663,17 @@ def test_ctes_in_order():
 
     sql = ibis.to_sql(expr, dialect="duckdb")
     assert sql.find('"first" AS (') < sql.find('"second" AS (')
+
+
+@pytest.mark.parametrize(
+    "accessor",
+    [
+        pytest.param(lambda s, f: getattr(s, f), id="getattr"),
+        pytest.param(lambda s, f: s[f], id="getitem"),
+    ],
+)
+def test_struct_field_simplified(snapshot, accessor):
+    s = ibis.struct({"i": 123, "s": "foo"})
+    for _i in range(5):
+        s = ibis.struct({"i": accessor(s, "i") + 1, "s": accessor(s, "s") + "bar"})
+    snapshot.assert_match(ibis.to_sql(s.i, dialect="duckdb"), "out.sql")

ibis/backends/tests/test_struct.py

@@ -78,14 +78,11 @@ def test_all_fields(struct, struct_df):
 
 
 @pytest.mark.notimpl(["postgres", "risingwave"])
-@pytest.mark.notyet(["datafusion"], raises=Exception, reason="unsupported syntax")
-@pytest.mark.parametrize("field", ["a", "b", "c"])
-def test_literal(backend, con, field):
-    query = _STRUCT_LITERAL[field]
-    dtype = query.type().to_pandas()
-    result = pd.Series([con.execute(query)], dtype=dtype)
+def test_literal(backend, con):
+    dtype = _STRUCT_LITERAL.type().to_pandas()
+    result = pd.Series([con.execute(_STRUCT_LITERAL)], dtype=dtype)
     result = result.replace({np.nan: None})
-    expected = pd.Series([_SIMPLE_DICT[field]])
+    expected = pd.Series([_SIMPLE_DICT])
     backend.assert_series_equal(result, expected.astype(dtype))
 
 

ibis/expr/types/structs.py

@@ -207,7 +207,18 @@ def __getitem__(self, name: str) -> ir.Value:
         """
         if name not in self.names:
             raise KeyError(name)
-        return ops.StructField(self, name).to_expr()
+
+        op = self.op()
+
+        # if the underlying operation is a simple struct column access, then
+        # just inline the underlying field access
+        if isinstance(op, ops.StructColumn):
+            return op.values[op.names.index(name)].to_expr()
+        # and then do the same if the underlying value is a field access
+        elif isinstance(op, ops.Literal):
+            return ops.Literal(op.value[name], dtype=self.fields[name]).to_expr()
+        else:
+            return ops.StructField(self, name).to_expr()
 
     def __setstate__(self, instance_dictionary):
         self.__dict__ = instance_dictionary

ibis/tests/expr/test_struct.py

@@ -31,15 +31,14 @@ def test_struct_operations():
         ]
     )
     expr = ibis.literal(value)
-    assert isinstance(expr, ir.StructValue)
-    assert isinstance(expr["b"], ir.ArrayValue)
-    assert isinstance(expr["a"].op(), ops.StructField)
+    assert isinstance(expr, ir.StructScalar)
+    assert isinstance(expr["b"], ir.ArrayScalar)
+    assert isinstance(expr["a"], ir.IntegerScalar)
 
 
 def test_struct_getattr():
     expr = ibis.struct({"a": 1, "b": 2})
     assert isinstance(expr.a, ir.IntegerValue)
-    assert expr.a.get_name() == "a"
     with pytest.raises(AttributeError, match="bad"):
         expr.bad  # # noqa: B018