refactor: Remove ibis usage to define new dataframe objects (#783)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -23,7 +23,6 @@
 import warnings
 
 import google.cloud.bigquery
-import ibis.expr.types as ibis_types
 import pandas
 import pyarrow as pa
 import pyarrow.feather as pa_feather
@@ -60,30 +59,6 @@ class ArrayValue:
 
     node: nodes.BigFrameNode
 
-    # DO NOT use, on deprecation path
-    @classmethod
-    def from_ibis(
-        cls,
-        session: Session,
-        table: ibis_types.Table,
-        columns: Sequence[ibis_types.Value],
-        hidden_ordering_columns: Sequence[ibis_types.Value],
-        ordering: orderings.ExpressionOrdering,
-    ):
-        import bigframes.core.compile.ibis_types
-
-        node = nodes.ReadGbqNode(
-            table=table,
-            table_session=session,
-            columns=tuple(
-                bigframes.core.compile.ibis_types.ibis_value_to_canonical_type(column)
-                for column in columns
-            ),
-            hidden_ordering_columns=tuple(hidden_ordering_columns),
-            ordering=ordering,
-        )
-        return cls(node)
-
     @classmethod
     def from_pyarrow(cls, arrow_table: pa.Table, session: Session):
         adapted_table = local_data.adapt_pa_table(arrow_table)

bigframes/core/blocks.py

@@ -2345,7 +2345,6 @@ def _get_rows_as_json_values(self) -> Block:
         # TODO(shobs): Replace direct SQL manipulation by structured expression
         # manipulation
         ordering_column_name = guid.generate_guid()
-        self.session._cache_with_offsets(self.expr)
         expr = self.expr.promote_offsets(ordering_column_name)
         expr_sql = self.session._to_sql(expr)
 
@@ -2415,17 +2414,31 @@ def _get_rows_as_json_values(self) -> Block:
 )
 SELECT {select_columns_csv} FROM T1
 """
-        ibis_table = self.session.ibis_client.sql(json_sql)
-        order_for_ibis_table = ordering.ExpressionOrdering.from_offset_col(
-            ordering_column_name
-        )
-        expr = core.ArrayValue.from_ibis(
-            self.session,
-            ibis_table,
-            [ibis_table[col] for col in select_columns if col != ordering_column_name],
-            hidden_ordering_columns=[ibis_table[ordering_column_name]],
-            ordering=order_for_ibis_table,
+        # The only ways this code is used is through df.apply(axis=1) cope path
+        destination, query_job = self.session._query_to_destination(
+            json_sql, index_cols=[ordering_column_name], api_name="apply"
+        )
+        if not destination:
+            raise ValueError(f"Query job {query_job} did not produce result table")
+
+        new_schema = (
+            self.expr.schema.select([*self.index_columns])
+            .append(
+                bf_schema.SchemaItem(
+                    row_json_column_name, bigframes.dtypes.STRING_DTYPE
+                )
+            )
+            .append(
+                bf_schema.SchemaItem(ordering_column_name, bigframes.dtypes.INT_DTYPE)
+            )
         )
+
+        expr = core.ArrayValue.from_table(
+            self.session.bqclient.get_table(destination),
+            schema=new_schema,
+            session=self.session,
+            offsets_col=ordering_column_name,
+        ).drop_columns([ordering_column_name])
         block = Block(
             expr,
             index_columns=self.index_columns,

bigframes/core/compile/compiler.py

@@ -223,22 +223,6 @@ def compile_read_table_ordered(node: nodes.ReadTableNode):
     )
 
 
-@_compile_node.register
-def compile_readgbq(node: nodes.ReadGbqNode, ordered: bool = True):
-    if ordered:
-        return compiled.OrderedIR(
-            node.table,
-            node.columns,
-            node.hidden_ordering_columns,
-            node.ordering,
-        )
-    else:
-        return compiled.UnorderedIR(
-            node.table,
-            node.columns,
-        )
-
-
 @_compile_node.register
 def compile_promote_offsets(node: nodes.PromoteOffsetsNode, ordered: bool = True):
     result = compile_ordered_ir(node.child).promote_offsets(node.col_id)

bigframes/core/nodes.py

@@ -34,8 +34,6 @@
 import bigframes.operations.aggregations as agg_ops
 
 if typing.TYPE_CHECKING:
-    import ibis.expr.types as ibis_types
-
     import bigframes.core.ordering as orderings
     import bigframes.session
 
@@ -302,54 +300,6 @@ def transform_children(
         return self
 
 
-# TODO: Refactor to take raw gbq object reference
-@dataclass(frozen=True)
-class ReadGbqNode(BigFrameNode):
-    table: ibis_types.Table = field()
-    table_session: bigframes.session.Session = field()
-    columns: Tuple[ibis_types.Value, ...] = field()
-    hidden_ordering_columns: Tuple[ibis_types.Value, ...] = field()
-    ordering: orderings.ExpressionOrdering = field()
-
-    @property
-    def session(self):
-        return self.table_session
-
-    def __hash__(self):
-        return self._node_hash
-
-    @property
-    def roots(self) -> typing.Set[BigFrameNode]:
-        return {self}
-
-    @functools.cached_property
-    def schema(self) -> schemata.ArraySchema:
-        from bigframes.core.compile.ibis_types import ibis_dtype_to_bigframes_dtype
-
-        items = tuple(
-            schemata.SchemaItem(
-                value.get_name(),
-                ibis_dtype_to_bigframes_dtype(value.type()),
-            )
-            for value in self.columns
-        )
-        return schemata.ArraySchema(items)
-
-    @functools.cached_property
-    def variables_introduced(self) -> int:
-        return len(self.columns) + len(self.hidden_ordering_columns)
-
-    @property
-    def relation_ops_created(self) -> int:
-        # Assume worst case, where readgbq actually has baked in analytic operation to generate index
-        return 2
-
-    def transform_children(
-        self, t: Callable[[BigFrameNode], BigFrameNode]
-    ) -> BigFrameNode:
-        return self
-
-
 ## Put ordering in here or just add order_by node above?
 @dataclass(frozen=True)
 class ReadTableNode(BigFrameNode):

tests/unit/resources.py

@@ -20,6 +20,7 @@
 import google.cloud.bigquery
 import ibis
 import pandas
+import pyarrow as pa
 import pytest
 
 import bigframes
@@ -130,18 +131,9 @@ def create_arrayvalue(
     df: pandas.DataFrame, total_ordering_columns: List[str]
 ) -> core.ArrayValue:
     session = create_pandas_session({"test_table": df})
-    ibis_table = session.ibis_client.table("test_table")
-    columns = tuple(ibis_table[key] for key in ibis_table.columns)
-    ordering = bigframes.core.ordering.ExpressionOrdering(
-        tuple(
-            [core.orderings.ascending_over(column) for column in total_ordering_columns]
-        ),
-        total_ordering_columns=frozenset(total_ordering_columns),
-    )
-    return core.ArrayValue.from_ibis(
+    return core.ArrayValue.from_pyarrow(
+        arrow_table=pa.Table.from_pandas(df, preserve_index=False),
         session=session,
-        table=ibis_table,
-        columns=columns,
-        hidden_ordering_columns=(),
-        ordering=ordering,
+    ).order_by(
+        [bigframes.core.ordering.ascending_over(col) for col in total_ordering_columns]
     )

tests/unit/test_core.py

@@ -15,45 +15,13 @@
 import ibis.expr.types as ibis_types
 import pandas
 
-import bigframes.core as core
 import bigframes.core.expression as ex
-import bigframes.core.ordering as order
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
 
 from . import resources
 
 
-def test_arrayvalue_constructor_from_ibis_table_adds_all_columns():
-    session = resources.create_pandas_session(
-        {
-            "test_table": pandas.DataFrame(
-                {
-                    "col1": [1, 2, 3],
-                    "not_included": [True, False, True],
-                    "col2": ["a", "b", "c"],
-                    "col3": [0.1, 0.2, 0.3],
-                }
-            )
-        }
-    )
-    ibis_table = session.ibis_client.table("test_table")
-    columns = (ibis_table["col1"], ibis_table["col2"], ibis_table["col3"])
-    ordering = order.ExpressionOrdering(
-        tuple([order.ascending_over("col1")]),
-        total_ordering_columns=frozenset(["col1"]),
-    )
-    actual = core.ArrayValue.from_ibis(
-        session=session,
-        table=ibis_table,
-        columns=columns,
-        ordering=ordering,
-        hidden_ordering_columns=(),
-    )
-    assert actual._compile_ordered()._table is ibis_table
-    assert len(actual.column_ids) == 3
-
-
 def test_arrayvalue_with_get_column_type():
     value = resources.create_arrayvalue(
         pandas.DataFrame(