fix: use object dtype for ARRAY columns in to_pandas() with pandas 1.x (#329)

* fix: use object dtype for ARRAY columns in `to_pandas()` with pandas 1.x * fix tests * simplify dtype logic

Author: tswast

bigframes/features.py

@@ -0,0 +1,42 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Tuple
+
+
+class PandasVersions:
+ """Version comparisons for pandas package"""
+
+ def __init__(self):
+ self._installed_version = None
+
+ @property
+ def installed_version(self) -> Tuple[str, ...]:
+ """pandas version"""
+ if self._installed_version is None:
+ import pandas
+
+ self._installed_version = tuple(pandas.__version__.split("."))
+ return self._installed_version
+
+ @property
+ def is_arrow_list_dtype_usable(self):
+ """True if pandas.ArrowDtype is usable."""
+ version = self.installed_version
+ return version[0] != "1"
+
+
+PANDAS_VERSIONS = PandasVersions()

bigframes/session/_io/pandas.py

@@ -19,8 +19,33 @@
 import pandas.arrays
 import pyarrow # type: ignore
 import pyarrow.compute # type: ignore
+import pyarrow.types # type: ignore
 
 import bigframes.constants
+import bigframes.features
+
+
+def _arrow_to_pandas_arrowdtype(
+ column: pyarrow.Array, dtype: pandas.ArrowDtype
+) -> pandas.Series:
+ if (
+ pyarrow.types.is_list(dtype.pyarrow_dtype)
+ and not bigframes.features.PANDAS_VERSIONS.is_arrow_list_dtype_usable
+ ):
+ # This version of pandas doesn't really support ArrowDtype
+ # well. See internal issue 321013333 where array type has
+ # several problems converting a string.
+ return pandas.Series(
+ column.to_pylist(), # type: ignore
+ dtype="object",
+ )
+
+ # Avoid conversion logic if we are backing the pandas Series by the
+ # arrow array.
+ return pandas.Series(
+ pandas.arrays.ArrowExtensionArray(column), # type: ignore
+ dtype=dtype,
+ )
 
 
 def arrow_to_pandas(
@@ -75,12 +100,7 @@ def arrow_to_pandas(
  )
  series = pandas.Series(pd_array, dtype=dtype)
  elif isinstance(dtype, pandas.ArrowDtype):
- # Avoid conversion logic if we are backing the pandas Series by the
- # arrow array.
- series = pandas.Series(
- pandas.arrays.ArrowExtensionArray(column), # type: ignore
- dtype=dtype,
- )
+ series = _arrow_to_pandas_arrowdtype(column, dtype)
  else:
  series = column.to_pandas(types_mapper=lambda _: dtype)