Skip to content

[python-package] Add support for PyArrow-backed Pandas DataFrames #7263

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.

Sign up for GitHub

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

Closed
maxzw wants to merge 2 commits into from
Closed

[python-package] Add support for PyArrow-backed Pandas DataFrames #7263

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2cc4190
Add support for PyArrow backed Pandas columns
maxzw May 7, 2026
035f680
Skip tests for old pandas versions
maxzw May 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 94 additions & 12 deletions python-package/lightgbm/basic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,11 +35,14 @@
arrow_is_floating,
arrow_is_integer,
concat,
pa_array,
pa_Array,
pa_chunked_array,
pa_ChunkedArray,
pa_compute,
pa_DataType,
pa_Table,
pd_ArrowDtype,
pd_CategoricalDtype,
pd_DataFrame,
pd_Series,
Expand Down Expand Up @@ -409,6 +412,23 @@ def _is_2d_collection(data: Any) -> bool:
return _is_numpy_2d_array(data) or _is_2d_list(data) or isinstance(data, pd_DataFrame)


def _is_pd_arrow_dtype(dtype: Any) -> bool:
"""Check whether the dtype is a pandas ArrowDtype."""
return pd_ArrowDtype is not None and isinstance(dtype, pd_ArrowDtype)


def _is_arrow_backed_pd_series(data: Any) -> bool:
"""Check whether data is a pandas Series backed by ArrowDtype."""
return isinstance(data, pd_Series) and _is_pd_arrow_dtype(data.dtype)


def _extract_arrow_array(data: Any) -> Any:
"""Extract PyArrow array from arrow-backed pandas Series, or return as-is if already an arrow array."""
if _is_arrow_backed_pd_series(data):
return data.array.__arrow_array__()
return data


def _is_pyarrow_array(data: Any) -> "TypeGuard[Union[pa_Array, pa_ChunkedArray]]":
"""Check whether data is a PyArrow array."""
return isinstance(data, (pa_Array, pa_ChunkedArray))
Expand Down Expand Up @@ -790,6 +810,55 @@ def _c_int_array(data: np.ndarray) -> Tuple[_ctypes_int_ptr, int, np.ndarray]:
return (ptr_data, type_data, data) # return `data` to avoid the temporary copy is freed


def _is_allowed_arrow_dtype(arrow_type: pa_DataType) -> bool:
"""Check if the Arrow type is one that C++ can handle (integer, floating point, or boolean)."""
return arrow_is_integer(arrow_type) or arrow_is_floating(arrow_type) or arrow_is_boolean(arrow_type)


def _check_for_bad_arrow_dtypes(table: pa_Table) -> None:
"""Raise if the Arrow table contains non-numeric types."""
bad_arrow_dtypes = [
f"{field.name}: {field.type}"
for field in table.schema
if not _is_allowed_arrow_dtype(field.type)
]
if bad_arrow_dtypes:
raise ValueError(
"Arrow columns must be int, float or bool.\n"
f"Fields with bad Arrow dtypes: {', '.join(bad_arrow_dtypes)}"
)


def _pandas_to_arrow_table(data: pd_DataFrame) -> pa_Table:
"""Convert a pandas DataFrame to an Arrow Table."""
if not (PYARROW_INSTALLED and CFFI_INSTALLED):
raise LightGBMError("Cannot init Dataset from Arrow without 'pyarrow' and 'cffi' installed.")
arrays = []
names = []
for col_name, col in data.items():
dtype = col.dtype
if _is_pd_arrow_dtype(dtype):
# arrow-backed columns can be converted zero-copy by extracting the underlying array
arrays.append(col.array.__arrow_array__())
elif hasattr(dtype, 'numpy_dtype'):
# nullable extension types (e.g. pd.Int64Dtype, pd.Float64Dtype, pd.BooleanDtype)
# use separate data+mask arrays; pass the Series directly so pyarrow can interpret the mask
arrays.append(pa_array(col, from_pandas=True))
elif _is_allowed_numpy_dtype(dtype.type):
# numpy-backed columns are a single contiguous buffer; .values is a zero-copy view
arrays.append(pa_array(col.values, from_pandas=True))
else:
raise ValueError(
f"Column '{col_name}' has unsupported dtype '{dtype}'. "
"Columns must have integer, floating point, or boolean dtypes."
)
names.append(str(col_name))
# column names are coerced to strings as required by Arrow
table = pa_Table.from_arrays(arrays, names=names)
_check_for_bad_arrow_dtypes(table)
return table


def _is_allowed_numpy_dtype(dtype: type) -> bool:
float128 = getattr(np, "float128", type(None))
return issubclass(dtype, (np.integer, np.floating, np.bool_)) and not issubclass(dtype, (np.timedelta64, float128))
Expand Down Expand Up @@ -829,7 +898,7 @@ def _data_from_pandas(
feature_name: _LGBM_FeatureNameConfiguration,
categorical_feature: _LGBM_CategoricalFeatureConfiguration,
pandas_categorical: Optional[List[List]],
) -> Tuple[np.ndarray, List[str], Union[List[str], List[int]], List[List]]:
) -> Tuple[Union[np.ndarray, pa_Table], List[str], Union[List[str], List[int]], List[List]]:
if len(data.shape) != 2 or data.shape[0] < 1:
raise ValueError("Input data must be 2 dimensional and non empty.")

Expand Down Expand Up @@ -859,6 +928,14 @@ def _data_from_pandas(
if categorical_feature == "auto":
categorical_feature = cat_cols_not_ordered

if any(_is_pd_arrow_dtype(dtype) for dtype in data.dtypes):
return (
_pandas_to_arrow_table(data),
feature_name,
categorical_feature,
pandas_categorical,
)

df_dtypes = [dtype.type for dtype in data.dtypes]
# so that the target dtype considers floats
df_dtypes.append(np.float32)
Expand Down Expand Up @@ -1718,9 +1795,7 @@ def __pred_for_pyarrow_table(
if not (PYARROW_INSTALLED and CFFI_INSTALLED):
raise LightGBMError("Cannot predict from Arrow without 'pyarrow' and 'cffi' installed.")

# Check that the input is valid: we only handle numbers (for now)
if not all(arrow_is_integer(t) or arrow_is_floating(t) or arrow_is_boolean(t) for t in table.schema.types):
raise ValueError("Arrow table may only have integer or floating point datatypes")
_check_for_bad_arrow_dtypes(table)

# Prepare prediction output array
n_preds = self.__get_num_preds(
Expand Down Expand Up @@ -2473,9 +2548,7 @@ def __init_from_pyarrow_table(
if not (PYARROW_INSTALLED and CFFI_INSTALLED):
raise LightGBMError("Cannot init Dataset from Arrow without 'pyarrow' and 'cffi' installed.")

# Check that the input is valid: we only handle numbers (for now)
if not all(arrow_is_integer(t) or arrow_is_floating(t) or arrow_is_boolean(t) for t in table.schema.types):
raise ValueError("Arrow table may only have integer or floating point datatypes")
_check_for_bad_arrow_dtypes(table)

# Export Arrow table to C
c_array = _export_arrow_to_c(table)
Expand Down Expand Up @@ -2804,6 +2877,8 @@ def set_field(
)
return self

data = _extract_arrow_array(data)

# If the data is a arrow data, we can just pass it to C
if _is_pyarrow_array(data) or _is_pyarrow_table(data):
# If a table is being passed, we concatenate the columns. This is only valid for
Expand Down Expand Up @@ -3076,6 +3151,7 @@ def set_label(self, label: Optional[_LGBM_LabelType]) -> "Dataset":
----------
label : list, numpy 1-D array, pandas Series / one-column DataFrame, pyarrow Array, pyarrow ChunkedArray or None
The label information to be set into Dataset.
Supports arrow-backed pandas Series (ArrowDtype).

Returns
-------
Expand All @@ -3088,7 +3164,7 @@ def set_label(self, label: Optional[_LGBM_LabelType]) -> "Dataset":
if len(label.columns) > 1:
raise ValueError("DataFrame for label cannot have multiple columns")
label_array = np.ravel(_pandas_to_numpy(label, target_dtype=np.float32))
elif _is_pyarrow_array(label):
elif _is_pyarrow_array(label) or _is_arrow_backed_pd_series(label):
label_array = label
else:
label_array = _list_to_1d_numpy(data=label, dtype=np.float32, name="label")
Expand All @@ -3106,6 +3182,7 @@ def set_weight(
----------
weight : list, numpy 1-D array, pandas Series, pyarrow Array, pyarrow ChunkedArray or None
Weight to be set for each data point. Weights should be non-negative.
Supports arrow-backed pandas Series (ArrowDtype).

Returns
-------
Expand All @@ -3114,6 +3191,7 @@ def set_weight(
"""
# Check if the weight contains values other than one
if weight is not None:
weight = _extract_arrow_array(weight)
if _is_pyarrow_array(weight):
# TODO: remove 'type: ignore[attr-defined]' when https://github.com/apache/arrow/issues/49831 is resolved.
if pa_compute.all(pa_compute.equal(weight, 1)).as_py(): # type: ignore[attr-defined]
Expand All @@ -3124,7 +3202,7 @@ def set_weight(

# Set field
if self._handle is not None and weight is not None:
if not _is_pyarrow_array(weight):
if not (_is_pyarrow_array(weight) or _is_arrow_backed_pd_series(weight)):
weight = _list_to_1d_numpy(data=weight, dtype=np.float32, name="weight")
self.set_field("weight", weight)
self.weight = self.get_field("weight") # original values can be modified at cpp side
Expand All @@ -3140,6 +3218,7 @@ def set_init_score(
----------
init_score : list, list of lists (for multi-class task), numpy array, pandas Series, pandas DataFrame (for multi-class task), pyarrow Array, pyarrow ChunkedArray, pyarrow Table (for multi-class task) or None
Init score for Booster.
Supports arrow-backed pandas Series and DataFrame (ArrowDtype).

Returns
-------
Expand All @@ -3166,6 +3245,7 @@ def set_group(
sum(group) = n_samples.
For example, if you have a 100-document dataset with ``group = [10, 20, 40, 10, 10, 10]``, that means that you have 6 groups,
where the first 10 records are in the first group, records 11-30 are in the second group, records 31-70 are in the third group, etc.
Supports arrow-backed pandas Series (ArrowDtype).

Returns
-------
Expand All @@ -3174,7 +3254,7 @@ def set_group(
"""
self.group = group
if self._handle is not None and group is not None:
if not _is_pyarrow_array(group):
if not (_is_pyarrow_array(group) or _is_arrow_backed_pd_series(group)):
group = _list_to_1d_numpy(data=group, dtype=np.int32, name="group")
self.set_field("group", group)
# original values can be modified at cpp side
Expand All @@ -3191,8 +3271,9 @@ def set_position(

Parameters
----------
position : numpy 1-D array, pandas Series or None, optional (default=None)
position : numpy 1-D array, pandas Series, pyarrow Array, pyarrow ChunkedArray or None, optional (default=None)
Position of items used in unbiased learning-to-rank task.
Supports arrow-backed pandas Series (ArrowDtype).

Returns
-------
Expand All @@ -3201,7 +3282,8 @@ def set_position(
"""
self.position = position
if self._handle is not None and position is not None:
position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position")
if not (_is_pyarrow_array(position) or _is_arrow_backed_pd_series(position)):
position = _list_to_1d_numpy(data=position, dtype=np.int32, name="position")
self.set_field("position", position)
return self

Expand Down
17 changes: 17 additions & 0 deletions python-package/lightgbm/compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -176,6 +176,10 @@ class _LGBMRegressorBase: # type: ignore
from pandas import CategoricalDtype as pd_CategoricalDtype
except ImportError:
from pandas.api.types import CategoricalDtype as pd_CategoricalDtype
try:
from pandas import ArrowDtype as pd_ArrowDtype
except ImportError:
pd_ArrowDtype = None
PANDAS_INSTALLED = True
except ImportError:
PANDAS_INSTALLED = False
Expand All @@ -198,6 +202,12 @@ class pd_CategoricalDtype: # type: ignore
def __init__(self, *args: Any, **kwargs: Any):
pass

class pd_ArrowDtype: # type: ignore
"""Dummy class for pandas.ArrowDtype."""

def __init__(self, *args: Any, **kwargs: Any):
pass

concat = None

"""matplotlib"""
Expand Down Expand Up @@ -285,6 +295,7 @@ def __init__(self, *args: Any, **kwargs: Any):
from pyarrow import Table as pa_Table
from pyarrow import array as pa_array
from pyarrow import chunked_array as pa_chunked_array
from pyarrow.lib import DataType as pa_DataType
from pyarrow.types import is_boolean as arrow_is_boolean
from pyarrow.types import is_floating as arrow_is_floating
from pyarrow.types import is_integer as arrow_is_integer
Expand All @@ -311,6 +322,12 @@ class pa_Table: # type: ignore
def __init__(self, *args: Any, **kwargs: Any):
pass

class pa_DataType: # type: ignore
"""Dummy class for pa.lib.DataType."""

def __init__(self, *args: Any, **kwargs: Any):
pass

class pa_compute: # type: ignore
"""Dummy class for pyarrow.compute module."""

Expand Down
Loading
Loading