fix schema loacing

Author: TrevorBergeron

bigframes/core/array_value.py

@@ -88,9 +88,9 @@ def from_range(cls, start, end, step):
     def from_table(
         cls,
         table: google.cloud.bigquery.Table,
-        schema: schemata.ArraySchema,
         session: Session,
         *,
+        columns: Optional[Sequence[str]] = None,
         predicate: Optional[str] = None,
         at_time: Optional[datetime.datetime] = None,
         primary_key: Sequence[str] = (),
@@ -100,7 +100,7 @@ def from_table(
         if offsets_col and primary_key:
             raise ValueError("must set at most one of 'offests', 'primary_key'")
         # define data source only for needed columns, this makes row-hashing cheaper
-        table_def = bq_data.GbqTable.from_table(table, columns=schema.names)
+        table_def = bq_data.GbqTable.from_table(table, columns=columns or ())
 
         # create ordering from info
         ordering = None
@@ -111,16 +111,17 @@ def from_table(
                 [ids.ColumnId(key_part) for key_part in primary_key]
             )
 
+        bf_schema = schemata.ArraySchema.from_bq_table(table, columns=columns)
         # Scan all columns by default, we define this list as it can be pruned while preserving source_def
         scan_list = nodes.ScanList(
             tuple(
                 nodes.ScanItem(ids.ColumnId(item.column), item.column)
-                for item in schema.items
+                for item in bf_schema.items
             )
         )
         source_def = bq_data.BigqueryDataSource(
             table=table_def,
-            schema=schema,
+            schema=bf_schema,
             at_time=at_time,
             sql_predicate=predicate,
             ordering=ordering,

bigframes/core/bq_data.py

@@ -84,7 +84,8 @@ class BigqueryDataSource:
     """
 
     def __post_init__(self):
-        assert [field.name for field in self.table.physical_schema] == list(
+        # not all columns need be in schema, eg so can exclude unsupported column types (eg RANGE)
+        assert set(field.name for field in self.table.physical_schema).issuperset(
             self.schema.names
         )
 

bigframes/core/schema.py

@@ -17,7 +17,7 @@
 from dataclasses import dataclass
 import functools
 import typing
-from typing import Dict, List
+from typing import Dict, List, Optional, Sequence
 
 import google.cloud.bigquery
 import pyarrow
@@ -44,21 +44,26 @@ def __iter__(self):
     def from_bq_table(
         cls,
         table: google.cloud.bigquery.Table,
-        column_type_overrides: typing.Optional[
+        column_type_overrides: Optional[
             typing.Dict[str, bigframes.dtypes.Dtype]
         ] = None,
+        columns: Optional[Sequence[str]] = None,
     ):
+        if not columns:
+            fields = table.schema
+        else:
+            lookup = {field.name: field for field in table.schema}
+            fields = [lookup[col] for col in columns]
+
         return ArraySchema.from_bq_schema(
-            table.schema, column_type_overrides=column_type_overrides
+            fields, column_type_overrides=column_type_overrides
         )
 
     @classmethod
     def from_bq_schema(
         cls,
         schema: List[google.cloud.bigquery.SchemaField],
-        column_type_overrides: typing.Optional[
-            Dict[str, bigframes.dtypes.Dtype]
-        ] = None,
+        column_type_overrides: Optional[Dict[str, bigframes.dtypes.Dtype]] = None,
     ):
         if column_type_overrides is None:
             column_type_overrides = {}
@@ -90,14 +95,16 @@ def to_bigquery(
             for item in self.items
         )
 
-    def to_pyarrow(self) -> pyarrow.Schema:
+    def to_pyarrow(self, use_storage_type: bool = False) -> pyarrow.Schema:
         fields = []
         for item in self.items:
             pa_type = bigframes.dtypes.bigframes_dtype_to_arrow_dtype(item.dtype)
             fields.append(
                 pyarrow.field(
                     item.column,
-                    pa_type,
+                    pa_type.storage_type
+                    if use_storage_type and isinstance(pa_type, pyarrow.ExtensionType)
+                    else pa_type,
                     nullable=not pyarrow.types.is_list(pa_type),
                 )
             )

bigframes/session/executor.py

@@ -100,7 +100,10 @@ def to_arrow_table(self) -> pyarrow.Table:
                 itertools.chain(peek_value, batches),  # reconstruct
             )
         else:
-            return self._schema.to_pyarrow().empty_table()
+            try:
+                return self._schema.to_pyarrow().empty_table()
+            except pa.ArrowNotImplementedError:
+                return self._schema.to_pyarrow(use_storage_type=True).empty_table()
 
     def to_pandas(self) -> pd.DataFrame:
         return io_pandas.arrow_to_pandas(self.to_arrow_table(), self._schema)

bigframes/session/loader.py

@@ -811,12 +811,10 @@ def read_gbq_table(
                     bigframes.core.events.ExecutionFinished(),
                 )
 
-        schema = schemata.ArraySchema.from_bq_table(table)
-        if not include_all_columns:
-            schema = schema.select(index_cols + columns)
+        selected_cols = None if include_all_columns else index_cols + columns
         array_value = core.ArrayValue.from_table(
             table,
-            schema=schema,
+            columns=selected_cols,
             predicate=filter_str,
             at_time=time_travel_timestamp if enable_snapshot else None,
             primary_key=primary_key,

tests/unit/core/rewrite/conftest.py

@@ -72,7 +72,6 @@ def leaf(fake_session, table):
     return core.ArrayValue.from_table(
         session=fake_session,
         table=table,
-        schema=bigframes.core.schema.ArraySchema.from_bq_table(table),
     ).node
 
 
@@ -81,5 +80,4 @@ def leaf_too(fake_session, table_too):
     return core.ArrayValue.from_table(
         session=fake_session,
         table=table_too,
-        schema=bigframes.core.schema.ArraySchema.from_bq_table(table_too),
     ).node

tests/unit/test_planner.py

@@ -39,7 +39,6 @@
 LEAF: core.ArrayValue = core.ArrayValue.from_table(
     session=FAKE_SESSION,
     table=TABLE,
-    schema=bigframes.core.schema.ArraySchema.from_bq_table(TABLE),
 )