-
-
Couldn't load subscription status.
- Fork 19.2k
ENH/TST: Add BaseInterfaceTests tests for ArrowExtensionArray #47377
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
2c984c2
a766ed1
3b01b77
e09ac0a
7ca44ae
2c06f26
607baed
5e7af38
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -31,6 +31,7 @@ | |
| ) | ||
| from pandas.core.dtypes.missing import isna | ||
|
|
||
| from pandas.core.arraylike import OpsMixin | ||
| from pandas.core.arrays.base import ExtensionArray | ||
| from pandas.core.indexers import ( | ||
| check_array_indexer, | ||
|
|
@@ -45,13 +46,22 @@ | |
| from pandas.core.arrays.arrow._arrow_utils import fallback_performancewarning | ||
| from pandas.core.arrays.arrow.dtype import ArrowDtype | ||
|
|
||
| ARROW_CMP_FUNCS = { | ||
| "eq": pc.equal, | ||
| "ne": pc.not_equal, | ||
| "lt": pc.less, | ||
| "gt": pc.greater, | ||
| "le": pc.less_equal, | ||
| "ge": pc.greater_equal, | ||
| } | ||
|
|
||
| if TYPE_CHECKING: | ||
| from pandas import Series | ||
|
|
||
| ArrowExtensionArrayT = TypeVar("ArrowExtensionArrayT", bound="ArrowExtensionArray") | ||
|
|
||
|
|
||
| class ArrowExtensionArray(ExtensionArray): | ||
| class ArrowExtensionArray(OpsMixin, ExtensionArray): | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Should I be using There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. i recommend OpsMixin |
||
| """ | ||
| Base class for ExtensionArray backed by Arrow ChunkedArray. | ||
| """ | ||
|
|
@@ -179,6 +189,34 @@ def __arrow_array__(self, type=None): | |
| """Convert myself to a pyarrow ChunkedArray.""" | ||
| return self._data | ||
|
|
||
| def _cmp_method(self, other, op): | ||
| from pandas.arrays import BooleanArray | ||
|
|
||
| pc_func = ARROW_CMP_FUNCS[op.__name__] | ||
| if isinstance(other, ArrowExtensionArray): | ||
| result = pc_func(self._data, other._data) | ||
| elif isinstance(other, (np.ndarray, list)): | ||
| result = pc_func(self._data, other) | ||
| elif is_scalar(other): | ||
| try: | ||
| result = pc_func(self._data, pa.scalar(other)) | ||
| except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid): | ||
| mask = isna(self) | isna(other) | ||
| valid = ~mask | ||
| result = np.zeros(len(self), dtype="bool") | ||
| result[valid] = op(np.array(self)[valid], other) | ||
| return BooleanArray(result, mask) | ||
| else: | ||
| return NotImplementedError( | ||
| f"{op.__name__} not implemented for {type(other)}" | ||
| ) | ||
|
|
||
| if pa_version_under2p0: | ||
| result = result.to_pandas().values | ||
| else: | ||
| result = result.to_numpy() | ||
| return BooleanArray._from_sequence(result) | ||
|
|
||
| def equals(self, other) -> bool: | ||
| if not isinstance(other, ArrowExtensionArray): | ||
| return False | ||
|
|
@@ -581,7 +619,7 @@ def _replace_with_indices( | |
| # fast path for a contiguous set of indices | ||
| arrays = [ | ||
| chunk[:start], | ||
| pa.array(value, type=chunk.type), | ||
| pa.array(value, type=chunk.type, from_pandas=True), | ||
| chunk[stop + 1 :], | ||
| ] | ||
| arrays = [arr for arr in arrays if len(arr)] | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
isn't there a min version on some of these?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks like they all have been around since pyarrow version 2 which I think is the version we require to use these functions
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Plus we catch the
except (pa.lib.ArrowNotImplementedError, pa.lib.ArrowInvalid)when calling these with fallback behavior.