[python-package] add PyArrow Table support to get_data and add_features_from methods #6892
+220
−2
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -46,6 +46,9 @@ | |
| pd_Series, | ||
| ) | ||
|
|
||
| if PYARROW_INSTALLED: | ||
| import pyarrow as pa | ||
|
|
||
| if TYPE_CHECKING: | ||
| from typing import Literal | ||
|
|
||
|
|
@@ -3293,6 +3296,8 @@ def get_data(self) -> Optional[_LGBM_TrainDataType]: | |
| self.data = self.data[self.used_indices, :] | ||
| elif isinstance(self.data, Sequence): | ||
| self.data = self.data[self.used_indices] | ||
| elif isinstance(self.data, pa_Table): | ||
| self.data = self.data.take(self.used_indices) | ||
| elif _is_list_of_sequences(self.data) and len(self.data) > 0: | ||
| self.data = np.array(list(self._yield_row_from_seqlist(self.data, self.used_indices))) | ||
| else: | ||
|
|
@@ -3480,6 +3485,22 @@ def add_features_from(self, other: "Dataset") -> "Dataset": | |
| elif isinstance(other.data, dt_DataTable): | ||
| _emit_datatable_deprecation_warning() | ||
| self.data = np.hstack((self.data, other.data.to_numpy())) | ||
| elif isinstance(other.data, pa_Table): | ||
| if not PYARROW_INSTALLED: | ||
| raise LightGBMError( | ||
| "Cannot add features to pyarrow.Table type of raw data " | ||
| "without pyarrow installed. " | ||
| "Install pyarrow and restart your session." | ||
| ) | ||
| else: | ||
|
||
| self.data = np.hstack( | ||
| ( | ||
| self.data, | ||
| np.column_stack( | ||
| [other.data.column(i).to_numpy() for i in range(len(other.data.column_names))] | ||
| ), | ||
| ) | ||
| ) | ||
| else: | ||
| self.data = None | ||
| elif isinstance(self.data, scipy.sparse.spmatrix): | ||
|
|
@@ -3491,6 +3512,23 @@ def add_features_from(self, other: "Dataset") -> "Dataset": | |
| elif isinstance(other.data, dt_DataTable): | ||
| _emit_datatable_deprecation_warning() | ||
| self.data = scipy.sparse.hstack((self.data, other.data.to_numpy()), format=sparse_format) | ||
| elif isinstance(other.data, pa_Table): | ||
| if not PYARROW_INSTALLED: | ||
| raise LightGBMError( | ||
| "Cannot add features to pyarrow.Table type of raw data " | ||
| "without pyarrow installed. " | ||
| "Install pyarrow and restart your session." | ||
| ) | ||
| else: | ||
| self.data = scipy.sparse.hstack( | ||
| ( | ||
| self.data, | ||
| np.column_stack( | ||
| [other.data.column(i).to_numpy() for i in range(len(other.data.column_names))] | ||
| ), | ||
| ), | ||
| format=sparse_format, | ||
| ) | ||
| else: | ||
| self.data = None | ||
| elif isinstance(self.data, pd_DataFrame): | ||
|
|
@@ -3509,6 +3547,27 @@ def add_features_from(self, other: "Dataset") -> "Dataset": | |
| elif isinstance(other.data, dt_DataTable): | ||
| _emit_datatable_deprecation_warning() | ||
| self.data = concat((self.data, pd_DataFrame(other.data.to_numpy())), axis=1, ignore_index=True) | ||
| elif isinstance(other.data, pa_Table): | ||
| if not PYARROW_INSTALLED: | ||
| raise LightGBMError( | ||
| "Cannot add features to pyarrow.Table type of raw data " | ||
| "without pyarrow installed. " | ||
| "Install pyarrow and restart your session." | ||
| ) | ||
| else: | ||
| self.data = concat( | ||
| ( | ||
| self.data, | ||
| pd_DataFrame( | ||
| { | ||
| other.data.column_names[i]: other.data.column(i).to_numpy() | ||
| for i in range(len(other.data.column_names)) | ||
| } | ||
| ), | ||
| ), | ||
| axis=1, | ||
| ignore_index=True, | ||
| ) | ||
| else: | ||
| self.data = None | ||
| elif isinstance(self.data, dt_DataTable): | ||
|
|
@@ -3521,6 +3580,79 @@ def add_features_from(self, other: "Dataset") -> "Dataset": | |
| self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data.values))) | ||
| elif isinstance(other.data, dt_DataTable): | ||
| self.data = dt_DataTable(np.hstack((self.data.to_numpy(), other.data.to_numpy()))) | ||
| elif isinstance(other.data, pa_Table): | ||
| if not PYARROW_INSTALLED: | ||
| raise LightGBMError( | ||
| "Cannot add features to pyarrow.Table type of raw data " | ||
| "without pyarrow installed. " | ||
| "Install pyarrow and restart your session." | ||
| ) | ||
| else: | ||
| self.data = dt_DataTable( | ||
| np.hstack( | ||
| ( | ||
| self.data.to_numpy(), | ||
| np.column_stack( | ||
| [other.data.column(i).to_numpy() for i in range(len(other.data.column_names))] | ||
| ), | ||
| ) | ||
| ) | ||
| ) | ||
| else: | ||
| self.data = None | ||
|
||
| elif isinstance(self.data, pa_Table): | ||
| if not PYARROW_INSTALLED: | ||
| raise LightGBMError( | ||
| "Cannot add features to pyarrow.Table type of raw data " | ||
| "without pyarrow installed. " | ||
| "Install pyarrow and restart your session." | ||
| ) | ||
| if isinstance(other.data, np.ndarray): | ||
| self.data = pa_Table.from_arrays( | ||
| [ | ||
| *self.data.columns, | ||
| *[pa.array(other.data[:, i]) for i in range(other.data.shape[1])], | ||
| ], | ||
| names=[ | ||
| *self.data.column_names, | ||
| *[f"D{len(self.data.column_names) + i + 1}" for i in range(other.data.shape[1])], | ||
| ], | ||
| ) | ||
| elif isinstance(other.data, scipy.sparse.spmatrix): | ||
| other_array = other.data.toarray() | ||
| self.data = pa_Table.from_arrays( | ||
| [ | ||
| *self.data.columns, | ||
| *[pa.array(other_array[:, i]) for i in range(other_array.shape[1])], | ||
| ], | ||
| names=[ | ||
| *self.data.column_names, | ||
| *[f"D{len(self.data.column_names) + i + 1}" for i in range(other_array.shape[1])], | ||
| ], | ||
| ) | ||
| elif isinstance(other.data, pd_DataFrame): | ||
| self.data = pa_Table.from_arrays( | ||
| [ | ||
| *self.data.columns, | ||
| *[pa.array(other.data.iloc[:, i].values) for i in range(len(other.data.columns))], | ||
| ], | ||
| names=[*self.data.column_names, *map(str, other.data.columns.tolist())], | ||
| ) | ||
| elif isinstance(other.data, dt_DataTable): | ||
| _emit_datatable_deprecation_warning() | ||
| other_array = other.data.to_numpy() | ||
| self.data = pa_Table.from_arrays( | ||
| [ | ||
| *self.data.columns, | ||
| *[pa.array(other_array[:, i]) for i in range(other_array.shape[1])], | ||
| ], | ||
| names=[*self.data.column_names, *other.data.names], | ||
| ) | ||
| elif isinstance(other.data, pa_Table): | ||
| self.data = pa_Table.from_arrays( | ||
| [*self.data.columns, *other.data.columns], | ||
| names=[*self.data.column_names, *other.data.column_names], | ||
| ) | ||
| else: | ||
| self.data = None | ||
| else: | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
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.
Thanks! This was definitely just something we'd missed.
Can you please add a test in https://github.com/microsoft/LightGBM/blob/master/tests/python_package_test/test_arrow.py just for this change to
get_data()? The other changes you made intest_basic.pydo not cover these changes.When you do that, please check that the content of
self.dataAND the returned value are correct (e.g., contain exactly the expected values and data types).If you'd like, I'd even support opening a new pull request that only has the
get_data()changes + test (and then making this PR only aboutadd_features_from()). Totally your choice, I want to be respectful of your time.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.
Thanks for the suggestion!
I added tests for both get_data() and add_features_from() directly in
test_arrow.pyas part of this PR.Please let me know if there’s anything else you’d like me to adjust!
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.
@jameslamb
I’ve opened a new pull request(#6911) that includes only the changes to get_data() along with the corresponding test. This should help keep things focused. I’d appreciate it if you could take a look when you have time.
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.
Thanks! I'll focus there.