fix extension type issue

Author: TrevorBergeron

bigframes/core/schema.py

@@ -95,16 +95,16 @@ def to_bigquery(
             for item in self.items
         )
 
-    def to_pyarrow(self, use_storage_type: bool = False) -> pyarrow.Schema:
+    def to_pyarrow(self, use_storage_types: bool = False) -> pyarrow.Schema:
         fields = []
         for item in self.items:
             pa_type = bigframes.dtypes.bigframes_dtype_to_arrow_dtype(item.dtype)
+            if use_storage_types:
+                pa_type = bigframes.dtypes.to_storage_type(pa_type)
             fields.append(
                 pyarrow.field(
                     item.column,
-                    pa_type.storage_type
-                    if use_storage_type and isinstance(pa_type, pyarrow.ExtensionType)
-                    else pa_type,
+                    type=pa_type,
                     nullable=not pyarrow.types.is_list(pa_type),
                 )
             )

bigframes/dtypes.py

@@ -553,6 +553,24 @@ def bigframes_dtype_to_arrow_dtype(
         )
 
 
+def to_storage_type(
+    arrow_type: pa.DataType,
+):
+    """Some pyarrow versions don't support extension types fully, such as for empty table generation."""
+    if isinstance(arrow_type, pa.ExtensionType):
+        return arrow_type.storage_type
+    if pa.types.is_list(arrow_type):
+        assert isinstance(arrow_type, pa.ListType)
+        return pa.list_(to_storage_type(arrow_type.value_type))
+    if pa.types.is_struct(arrow_type):
+        assert isinstance(arrow_type, pa.StructType)
+        return pa.struct(
+            field.with_type(to_storage_type(field.type))
+            for field in bigframes.core.backports.pyarrow_struct_type_fields(arrow_type)
+        )
+    return arrow_type
+
+
 def arrow_type_to_literal(
     arrow_type: pa.DataType,
 ) -> Any:

bigframes/session/executor.py

@@ -30,6 +30,7 @@
 import bigframes.core
 from bigframes.core import bq_data, local_data, pyarrow_utils
 import bigframes.core.schema
+import bigframes.dtypes
 import bigframes.session._io.pandas as io_pandas
 import bigframes.session.execution_spec as ex_spec
 
@@ -103,7 +104,8 @@ def to_arrow_table(self) -> pyarrow.Table:
             try:
                 return self._schema.to_pyarrow().empty_table()
             except pa.ArrowNotImplementedError:
-                return self._schema.to_pyarrow(use_storage_type=True).empty_table()
+                # Bug with some pyarrow versions, empty_table only supports base storage types, not extension types.
+                return self._schema.to_pyarrow(use_storage_types=True).empty_table()
 
     def to_pandas(self) -> pd.DataFrame:
         return io_pandas.arrow_to_pandas(self.to_arrow_table(), self._schema)