|
4 | 4 | from typing import (
|
5 | 5 | TYPE_CHECKING,
|
6 | 6 | Any,
|
7 |
| - Iterator, |
8 | 7 | Literal,
|
9 | 8 | Sequence,
|
10 | 9 | TypeVar,
|
@@ -855,65 +854,33 @@ def _set_via_chunk_iteration(
|
855 | 854 | """
|
856 | 855 | Loop through the array chunks and set the new values while
|
857 | 856 | leaving the chunking layout unchanged.
|
858 |
| - """ |
859 |
| - chunk_indices = self._indices_to_chunk_indices(indices) |
860 |
| - new_data = list(self._data.iterchunks()) |
861 |
| - |
862 |
| - for i, c_ind in enumerate(chunk_indices): |
863 |
| - n = len(c_ind) |
864 |
| - if n == 0: |
865 |
| - continue |
866 |
| - c_value, value = value[:n], value[n:] |
867 |
| - new_data[i] = self._replace_with_indices(new_data[i], c_ind, c_value) |
868 |
| - |
869 |
| - return pa.chunked_array(new_data) |
870 |
| - |
871 |
| - def _indices_to_chunk_indices( |
872 |
| - self, indices: npt.NDArray[np.intp] |
873 |
| - ) -> Iterator[npt.NDArray[np.intp]]: |
874 |
| - """ |
875 |
| - Convert *sorted* indices for self into a list of ndarrays |
876 |
| - each containing the indices *within* each chunk of the |
877 |
| - underlying ChunkedArray. |
878 | 857 |
|
879 | 858 | Parameters
|
880 | 859 | ----------
|
881 | 860 | indices : npt.NDArray[np.intp]
|
882 | 861 | Position indices for the underlying ChunkedArray.
|
883 | 862 |
|
884 |
| - Returns |
885 |
| - ------- |
886 |
| - Generator yielding positional indices for each chunk |
| 863 | + value : ExtensionDtype.type, Sequence[ExtensionDtype.type], or object |
| 864 | + value or values to be set of ``key``. |
887 | 865 |
|
888 | 866 | Notes
|
889 | 867 | -----
|
890 | 868 | Assumes that indices is sorted. Caller is responsible for sorting.
|
891 | 869 | """
|
892 |
| - for start, stop in self._chunk_positional_ranges(): |
| 870 | + new_data = [] |
| 871 | + stop = 0 |
| 872 | + for chunk in self._data.iterchunks(): |
| 873 | + start, stop = stop, stop + len(chunk) |
893 | 874 | if len(indices) == 0 or stop <= indices[0]:
|
894 |
| - yield np.array([], dtype=np.intp) |
| 875 | + new_data.append(chunk) |
895 | 876 | else:
|
896 | 877 | n = int(np.searchsorted(indices, stop, side="left"))
|
897 | 878 | c_ind = indices[:n] - start
|
898 | 879 | indices = indices[n:]
|
899 |
| - yield c_ind |
900 |
| - |
901 |
| - def _chunk_positional_ranges(self) -> tuple[tuple[int, int], ...]: |
902 |
| - """ |
903 |
| - Return a tuple of tuples each containing the left (inclusive) |
904 |
| - and right (exclusive) positional bounds of each chunk's values |
905 |
| - within the underlying ChunkedArray. |
906 |
| -
|
907 |
| - Returns |
908 |
| - ------- |
909 |
| - tuple[tuple] |
910 |
| - """ |
911 |
| - ranges = [] |
912 |
| - stop = 0 |
913 |
| - for c in self._data.iterchunks(): |
914 |
| - start, stop = stop, stop + len(c) |
915 |
| - ranges.append((start, stop)) |
916 |
| - return tuple(ranges) |
| 880 | + n = len(c_ind) |
| 881 | + c_value, value = value[:n], value[n:] |
| 882 | + new_data.append(self._replace_with_indices(chunk, c_ind, c_value)) |
| 883 | + return pa.chunked_array(new_data) |
917 | 884 |
|
918 | 885 | @classmethod
|
919 | 886 | def _replace_with_indices(
|
|
0 commit comments