FEAT-#7203: Make sure modin works correctly with pandas, which uses pyarrow as a backend

modin/core/dataframe/algebra/binary.py

@@ -13,8 +13,10 @@
 
 """Module houses builder class for Binary operator."""
 
+from __future__ import annotations
+
 import warnings
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import numpy as np
 import pandas
@@ -24,9 +26,12 @@
 
 from .operator import Operator
 
+if TYPE_CHECKING:
+ from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
+
 
 def maybe_compute_dtypes_common_cast(
- first,
+ first: PandasQueryCompiler,
  second,
  trigger_computations=False,
  axis=0,
@@ -80,6 +85,7 @@ def maybe_compute_dtypes_common_cast(
  # belong to the intersection, these will be NaN columns in the result
  mismatch_columns = columns_first ^ columns_second
  elif isinstance(second, dict):
+ # TODO: pyarrow backend
  dtypes_second = {
  key: pandas.api.types.pandas_dtype(type(value))
  for key, value in second.items()
@@ -92,6 +98,7 @@ def maybe_compute_dtypes_common_cast(
  mismatch_columns = columns_first.difference(columns_second)
  else:
  if isinstance(second, (list, tuple)):
+ # TODO: pyarrow backend
  second_dtypes_list = (
  [pandas.api.types.pandas_dtype(type(value)) for value in second]
  if axis == 1
@@ -100,6 +107,7 @@ def maybe_compute_dtypes_common_cast(
  else [np.array(second).dtype] * len(dtypes_first)
  )
  elif is_scalar(second) or isinstance(second, np.ndarray):
+ # TODO: pyarrow backend
  try:
  dtype = getattr(second, "dtype", None) or pandas.api.types.pandas_dtype(
  type(second)
@@ -125,6 +133,7 @@ def maybe_compute_dtypes_common_cast(
  mismatch_columns = []
 
  # If at least one column doesn't match, the result of the non matching column would be nan.
+ # TODO: pyarrow backend
  nan_dtype = pandas.api.types.pandas_dtype(type(np.nan))
  dtypes = None
  if func is not None:
@@ -168,7 +177,7 @@ def maybe_compute_dtypes_common_cast(
 
 
 def maybe_build_dtypes_series(
- first, second, dtype, trigger_computations=False
+ first: PandasQueryCompiler, second, dtype, trigger_computations=False
 ) -> Optional[pandas.Series]:
  """
  Build a ``pandas.Series`` describing dtypes of the result of a binary operation.
@@ -179,7 +188,7 @@ def maybe_build_dtypes_series(
  First operand for which the binary operation would be performed later.
  second : PandasQueryCompiler, list-like or scalar
  Second operand for which the binary operation would be performed later.
- dtype : np.dtype
+ dtype : pandas supported dtype
  Dtype of the result.
  trigger_computations : bool, default: False
  Whether to trigger computation of the lazy metadata for `first` and `second`.
@@ -217,8 +226,13 @@ def maybe_build_dtypes_series(
 
 
 def try_compute_new_dtypes(
- first, second, infer_dtypes=None, result_dtype=None, axis=0, func=None
-):
+ first: PandasQueryCompiler,
+ second,
+ infer_dtypes=None,
+ result_dtype=None,
+ axis=0,
+ func=None,
+) -> Optional[pandas.Series]:
  """
  Precompute resulting dtypes of the binary operation if possible.
 
@@ -235,7 +249,7 @@ def try_compute_new_dtypes(
  infer_dtypes : {"common_cast", "try_sample", "bool", None}, default: None
  How dtypes should be infered (see ``Binary.register`` doc for more info).
  result_dtype : np.dtype, optional
- NumPy dtype of the result. If not specified it will be inferred from the `infer_dtypes` parameter.
+ NumPy dtype of the result. If not specified it will be inferred from the `infer_dtypes` parameter. Only NumPy?
  axis : int, default: 0
  Axis to perform the binary operation along.
  func : callable(pandas.DataFrame, pandas.DataFrame) -> pandas.DataFrame, optional
@@ -250,10 +264,19 @@ def try_compute_new_dtypes(
 
  try:
  if infer_dtypes == "bool" or is_bool_dtype(result_dtype):
- # FIXME: https://github.com/modin-project/modin/issues/7203
- # can be `pandas.api.types.pandas_dtype("bool[pyarrow]")` depending on the data
+ # dataframe can contain types of different backends at the same time, for example:
+ # (Pdb) (pandas.DataFrame([[1,2,3], [4,5,6]]).astype({0: "int64[pyarrow]"}) > 4).dtypes
+ # 0 bool[pyarrow]
+ # 1 bool
+ # 2 bool
+ # dtype: object
+ backend = ""
+ if any("pyarrow" in str(x) for x in first.dtypes) or any(
+ "pyarrow" in str(x) for x in second.dtypes
+ ):
+ backend = "[pyarrow]"
  dtypes = maybe_build_dtypes_series(
- first, second, dtype=pandas.api.types.pandas_dtype(bool)
+ first, second, dtype=pandas.api.types.pandas_dtype(f"bool{backend}")
  )
  elif infer_dtypes == "common_cast":
  dtypes = maybe_compute_dtypes_common_cast(

modin/core/dataframe/algebra/map.py

@@ -13,8 +13,15 @@
 
 """Module houses builder class for Map operator."""
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from .operator import Operator
 
+if TYPE_CHECKING:
+ from modin.core.storage_formats.pandas.query_compiler import PandasQueryCompiler
+
 
 class Map(Operator):
  """Builder class for Map operator."""
@@ -41,7 +48,7 @@ def register(cls, function, *call_args, **call_kwds):
  Function that takes query compiler and executes map function.
  """
 
- def caller(query_compiler, *args, **kwargs):
+ def caller(query_compiler: PandasQueryCompiler, *args, **kwargs):
  """Execute Map function against passed query compiler."""
  shape_hint = call_kwds.pop("shape_hint", None) or query_compiler._shape_hint
  return query_compiler.__constructor__(

modin/core/dataframe/algebra/tree_reduce.py

@@ -35,7 +35,7 @@ def register(
  axis : int, optional
  Specifies axis to apply function along.
  compute_dtypes : callable(pandas.Series, *func_args, **func_kwargs) -> np.dtype, optional
- Callable for computing dtypes.
+ Callable for computing dtypes. Only NumPy?
 
  Returns
  -------

modin/core/dataframe/base/dataframe/utils.py

@@ -21,7 +21,6 @@
 from enum import Enum
 from typing import Dict, List, Sequence, Tuple, cast
 
-import numpy as np
 import pandas
 from pandas._typing import IndexLabel
 from pandas.api.types import is_scalar
@@ -170,7 +169,7 @@ def is_trivial_index(index: pandas.Index) -> bool:
  return True
  if isinstance(index, pandas.RangeIndex):
  return index.start == 0 and index.step == 1
- if not (isinstance(index, pandas.Index) and index.dtype == np.int64):
+ if not (isinstance(index, pandas.Index) and index.dtype == "int64"):
  return False
  return (
  index.is_monotonic_increasing