lazy import for now

Author: jorisvandenbossche

pandas/core/arrays/_arrow_utils.py

@@ -1,20 +1,12 @@
 from distutils.version import LooseVersion
+import json
 
 import numpy as np
+import pyarrow
 
-try:
-    import pyarrow
+from pandas.core.arrays.interval import _VALID_CLOSED
 
-    _PYARROW_INSTALLED = True
-except ImportError:
-    _PYARROW_INSTALLED = False
-    pyarrow = None
-
-
-if _PYARROW_INSTALLED:
-    _pyarrow_version_ge_015 = LooseVersion(pyarrow.__version__) >= LooseVersion("0.15")
-else:
-    _pyarrow_version_ge_015 = False
+_pyarrow_version_ge_015 = LooseVersion(pyarrow.__version__) >= LooseVersion("0.15")
 
 
 def pyarrow_array_to_numpy_and_mask(arr, dtype):
@@ -44,3 +36,89 @@ def pyarrow_array_to_numpy_and_mask(arr, dtype):
     else:
         mask = np.ones(len(arr), dtype=bool)
     return data, mask
+
+
+if _pyarrow_version_ge_015:
+    # the pyarrow extension types are only available for pyarrow 0.15+
+
+    class ArrowPeriodType(pyarrow.ExtensionType):
+        def __init__(self, freq):
+            # attributes need to be set first before calling
+            # super init (as that calls serialize)
+            self._freq = freq
+            pyarrow.ExtensionType.__init__(self, pyarrow.int64(), "pandas.period")
+
+        @property
+        def freq(self):
+            return self._freq
+
+        def __arrow_ext_serialize__(self):
+            metadata = {"freq": self.freq}
+            return json.dumps(metadata).encode()
+
+        @classmethod
+        def __arrow_ext_deserialize__(cls, storage_type, serialized):
+            metadata = json.loads(serialized.decode())
+            return ArrowPeriodType(metadata["freq"])
+
+        def __eq__(self, other):
+            if isinstance(other, pyarrow.BaseExtensionType):
+                return type(self) == type(other) and self.freq == other.freq
+            else:
+                return NotImplemented
+
+        def __hash__(self):
+            return hash((str(self), self.freq))
+
+    # register the type with a dummy instance
+    _period_type = ArrowPeriodType("D")
+    pyarrow.register_extension_type(_period_type)
+
+    class ArrowIntervalType(pyarrow.ExtensionType):
+        def __init__(self, subtype, closed):
+            # attributes need to be set first before calling
+            # super init (as that calls serialize)
+            assert closed in _VALID_CLOSED
+            self._closed = closed
+            if not isinstance(subtype, pyarrow.DataType):
+                subtype = pyarrow.type_for_alias(str(subtype))
+            self._subtype = subtype
+
+            storage_type = pyarrow.struct([("left", subtype), ("right", subtype)])
+            pyarrow.ExtensionType.__init__(self, storage_type, "pandas.interval")
+
+        @property
+        def subtype(self):
+            return self._subtype
+
+        @property
+        def closed(self):
+            return self._closed
+
+        def __arrow_ext_serialize__(self):
+            metadata = {"subtype": str(self.subtype), "closed": self.closed}
+            return json.dumps(metadata).encode()
+
+        @classmethod
+        def __arrow_ext_deserialize__(cls, storage_type, serialized):
+            metadata = json.loads(serialized.decode())
+            subtype = pyarrow.type_for_alias(metadata["subtype"])
+            closed = metadata["closed"]
+            return ArrowIntervalType(subtype, closed)
+
+        def __eq__(self, other):
+            if isinstance(other, pyarrow.BaseExtensionType):
+                return (
+                    type(self) == type(other)
+                    and self.subtype == other.subtype
+                    and self.closed == other.closed
+                )
+            else:
+                return NotImplemented
+
+        def __hash__(self):
+            return hash((str(self), str(self.subtype), self.closed))
+
+    # register the type with a dummy instance
+    _interval_type = ArrowIntervalType(pyarrow.int64(), "left")
+    pyarrow.register_extension_type(_interval_type)

pandas/core/arrays/integer.py

@@ -28,7 +28,6 @@
 from pandas.core import nanops, ops
 from pandas.core.algorithms import take
 from pandas.core.arrays import ExtensionArray, ExtensionOpsMixin
-from pandas.core.arrays._arrow_utils import pyarrow_array_to_numpy_and_mask
 import pandas.core.common as com
 from pandas.core.indexers import check_bool_array_indexer
 from pandas.core.ops import invalid_comparison
@@ -95,6 +94,7 @@ def construct_array_type(cls):
     def __from_arrow__(self, array):
         """Construct IntegerArray from passed pyarrow Array/ChunkedArray"""
         import pyarrow
+        from pandas.core.arrays._arrow_utils import pyarrow_array_to_numpy_and_mask
 
         if isinstance(array, pyarrow.Array):
             chunks = [array]

pandas/core/arrays/interval.py

@@ -1,4 +1,3 @@
-import json
 from operator import le, lt
 import textwrap
 
@@ -37,7 +36,6 @@
 from pandas.core.dtypes.missing import isna, notna
 
 from pandas.core.algorithms import take, value_counts
-from pandas.core.arrays._arrow_utils import _PYARROW_INSTALLED, _pyarrow_version_ge_015
 from pandas.core.arrays.base import ExtensionArray, _extension_array_shared_docs
 from pandas.core.arrays.categorical import Categorical
 import pandas.core.common as com
@@ -1088,6 +1086,7 @@ def __arrow_array__(self, type=None):
         Convert myself into a pyarrow Array.
         """
         import pyarrow
+        from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
         try:
             subtype = pyarrow.from_numpy_dtype(self.dtype.subtype)
@@ -1326,56 +1325,3 @@ def maybe_convert_platform_interval(values):
         values = np.asarray(values)
 
     return maybe_convert_platform(values)
-
-
-if _PYARROW_INSTALLED and _pyarrow_version_ge_015:
-    import pyarrow
-
-    class ArrowIntervalType(pyarrow.ExtensionType):
-        def __init__(self, subtype, closed):
-            # attributes need to be set first before calling
-            # super init (as that calls serialize)
-            assert closed in _VALID_CLOSED
-            self._closed = closed
-            if not isinstance(subtype, pyarrow.DataType):
-                subtype = pyarrow.type_for_alias(str(subtype))
-            self._subtype = subtype
-
-            storage_type = pyarrow.struct([("left", subtype), ("right", subtype)])
-            pyarrow.ExtensionType.__init__(self, storage_type, "pandas.interval")
-
-        @property
-        def subtype(self):
-            return self._subtype
-
-        @property
-        def closed(self):
-            return self._closed
-
-        def __arrow_ext_serialize__(self):
-            metadata = {"subtype": str(self.subtype), "closed": self.closed}
-            return json.dumps(metadata).encode()
-
-        @classmethod
-        def __arrow_ext_deserialize__(cls, storage_type, serialized):
-            metadata = json.loads(serialized.decode())
-            subtype = pyarrow.type_for_alias(metadata["subtype"])
-            closed = metadata["closed"]
-            return ArrowIntervalType(subtype, closed)
-
-        def __eq__(self, other):
-            if isinstance(other, pyarrow.BaseExtensionType):
-                return (
-                    type(self) == type(other)
-                    and self.subtype == other.subtype
-                    and self.closed == other.closed
-                )
-            else:
-                return NotImplemented
-
-        def __hash__(self):
-            return hash((str(self), str(self.subtype), self.closed))
-
-    # register the type with a dummy instance
-    _interval_type = ArrowIntervalType(pyarrow.int64(), "left")
-    pyarrow.register_extension_type(_interval_type)

pandas/core/arrays/period.py

@@ -1,5 +1,4 @@
 from datetime import timedelta
-import json
 import operator
 from typing import Any, Callable, List, Optional, Sequence, Union
 
@@ -42,7 +41,6 @@
 
 import pandas.core.algorithms as algos
 from pandas.core.arrays import datetimelike as dtl
-from pandas.core.arrays._arrow_utils import _PYARROW_INSTALLED, _pyarrow_version_ge_015
 import pandas.core.common as com
 
 from pandas.tseries import frequencies
@@ -290,6 +288,7 @@ def __arrow_array__(self, type=None):
         Convert myself into a pyarrow Array.
         """
         import pyarrow
+        from pandas.core.arrays._arrow_utils import ArrowPeriodType
 
         if type is not None:
             if pyarrow.types.is_integer(type):
@@ -1053,40 +1052,3 @@ def _make_field_arrays(*fields):
     ]
 
     return arrays
-
-
-if _PYARROW_INSTALLED and _pyarrow_version_ge_015:
-    import pyarrow
-
-    class ArrowPeriodType(pyarrow.ExtensionType):
-        def __init__(self, freq):
-            # attributes need to be set first before calling
-            # super init (as that calls serialize)
-            self._freq = freq
-            pyarrow.ExtensionType.__init__(self, pyarrow.int64(), "pandas.period")
-
-        @property
-        def freq(self):
-            return self._freq
-
-        def __arrow_ext_serialize__(self):
-            metadata = {"freq": self.freq}
-            return json.dumps(metadata).encode()
-
-        @classmethod
-        def __arrow_ext_deserialize__(cls, storage_type, serialized):
-            metadata = json.loads(serialized.decode())
-            return ArrowPeriodType(metadata["freq"])
-
-        def __eq__(self, other):
-            if isinstance(other, pyarrow.BaseExtensionType):
-                return type(self) == type(other) and self.freq == other.freq
-            else:
-                return NotImplemented
-
-        def __hash__(self):
-            return hash((str(self), self.freq))
-
-    # register the type with a dummy instance
-    _period_type = ArrowPeriodType("D")
-    pyarrow.register_extension_type(_period_type)

pandas/io/parquet.py

@@ -76,6 +76,9 @@ def __init__(self):
         )
         import pyarrow.parquet
 
+        # import utils to register the pyarrow extension types
+        import pandas.core.arrays._arrow_utils  # noqa
+
         self.api = pyarrow
 
     def write(

pandas/tests/arrays/interval/test_interval.py

@@ -117,7 +117,7 @@ def test_repr():
 @pyarrow_skip
 def test_arrow_extension_type():
     import pyarrow as pa
-    from pandas.core.arrays.interval import ArrowIntervalType
+    from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
     p1 = ArrowIntervalType(pa.int64(), "left")
     p2 = ArrowIntervalType(pa.int64(), "left")
@@ -133,7 +133,7 @@ def test_arrow_extension_type():
 @pyarrow_skip
 def test_arrow_array():
     import pyarrow as pa
-    from pandas.core.arrays.interval import ArrowIntervalType
+    from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
     intervals = pd.interval_range(1, 5, freq=1).array
 
@@ -162,7 +162,7 @@ def test_arrow_array():
 @pyarrow_skip
 def test_arrow_array_missing():
     import pyarrow as pa
-    from pandas.core.arrays.interval import ArrowIntervalType
+    from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
     arr = IntervalArray.from_breaks([0, 1, 2, 3])
     arr[1] = None
@@ -196,7 +196,7 @@ def test_arrow_array_missing():
 )
 def test_arrow_table_roundtrip(breaks):
     import pyarrow as pa
-    from pandas.core.arrays.interval import ArrowIntervalType
+    from pandas.core.arrays._arrow_utils import ArrowIntervalType
 
     arr = IntervalArray.from_breaks(breaks)
     arr[1] = None

pandas/tests/arrays/test_period.py

@@ -334,7 +334,7 @@ def test_min_max_empty(self, skipna):
 
 @pyarrow_skip
 def test_arrow_extension_type():
-    from pandas.core.arrays.period import ArrowPeriodType
+    from pandas.core.arrays._arrow_utils import ArrowPeriodType
 
     p1 = ArrowPeriodType("D")
     p2 = ArrowPeriodType("D")
@@ -357,7 +357,7 @@ def test_arrow_extension_type():
 )
 def test_arrow_array(data, freq):
     import pyarrow as pa
-    from pandas.core.arrays.period import ArrowPeriodType
+    from pandas.core.arrays._arrow_utils import ArrowPeriodType
 
     periods = period_array(data, freq=freq)
     result = pa.array(periods)
@@ -381,7 +381,7 @@ def test_arrow_array(data, freq):
 @pyarrow_skip
 def test_arrow_array_missing():
     import pyarrow as pa
-    from pandas.core.arrays.period import ArrowPeriodType
+    from pandas.core.arrays._arrow_utils import ArrowPeriodType
 
     arr = PeriodArray([1, 2, 3], freq="D")
     arr[1] = pd.NaT
@@ -396,7 +396,7 @@ def test_arrow_array_missing():
 @pyarrow_skip
 def test_arrow_table_roundtrip():
     import pyarrow as pa
-    from pandas.core.arrays.period import ArrowPeriodType
+    from pandas.core.arrays._arrow_utils import ArrowPeriodType
 
     arr = PeriodArray([1, 2, 3], freq="D")
     arr[1] = pd.NaT