| 929 | namespace internal { |
| 930 | |
| 931 | std::vector<ArrayVector> RechunkArraysConsistently( |
| 932 | const std::vector<ArrayVector>& groups) { |
| 933 | if (groups.size() <= 1) { |
| 934 | return groups; |
| 935 | } |
| 936 | int64_t total_length = 0; |
| 937 | for (const auto& array : groups.front()) { |
| 938 | total_length += array->length(); |
| 939 | } |
| 940 | #ifndef NDEBUG |
| 941 | for (const auto& group : groups) { |
| 942 | int64_t group_length = 0; |
| 943 | for (const auto& array : group) { |
| 944 | group_length += array->length(); |
| 945 | } |
| 946 | DCHECK_EQ(group_length, total_length) |
| 947 | << "Array groups should have the same total number of elements"; |
| 948 | } |
| 949 | #endif |
| 950 | if (total_length == 0) { |
| 951 | return groups; |
| 952 | } |
| 953 | |
| 954 | // Set up result vectors |
| 955 | std::vector<ArrayVector> rechunked_groups(groups.size()); |
| 956 | |
| 957 | // Set up progress counters |
| 958 | std::vector<ArrayVector::const_iterator> current_arrays; |
| 959 | std::vector<int64_t> array_offsets; |
| 960 | for (const auto& group : groups) { |
| 961 | current_arrays.emplace_back(group.cbegin()); |
| 962 | array_offsets.emplace_back(0); |
| 963 | } |
| 964 | |
| 965 | // Scan all array vectors at once, rechunking along the way |
| 966 | int64_t start = 0; |
| 967 | while (start < total_length) { |
| 968 | // First compute max possible length for next chunk |
| 969 | int64_t chunk_length = std::numeric_limits<int64_t>::max(); |
| 970 | for (size_t i = 0; i < groups.size(); i++) { |
| 971 | auto& arr_it = current_arrays[i]; |
| 972 | auto& offset = array_offsets[i]; |
| 973 | // Skip any done arrays (including 0-length arrays) |
| 974 | while (offset == (*arr_it)->length()) { |
| 975 | ++arr_it; |
| 976 | offset = 0; |
| 977 | } |
| 978 | const auto& array = *arr_it; |
| 979 | DCHECK_GT(array->length(), offset); |
| 980 | chunk_length = std::min(chunk_length, array->length() - offset); |
| 981 | } |
| 982 | DCHECK_GT(chunk_length, 0); |
| 983 | |
| 984 | // Then slice all arrays along this chunk size |
| 985 | for (size_t i = 0; i < groups.size(); i++) { |
| 986 | const auto& array = *current_arrays[i]; |
| 987 | auto& offset = array_offsets[i]; |
| 988 | if (offset == 0 && array->length() == chunk_length) { |