Created
January 24, 2024 15:30
-
-
Save amol-/b6cd98e62b09dce56f9ab29c65821bf2 to your computer and use it in GitHub Desktop.
This file contains 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
std::vector<Int64Builder> builders(num_chunks); | |
ChunkResolver index_resolver(values.chunks()); | |
for (int64_t requested_index = 0; requested_index < num_indices; ++requested_index) { | |
uint64_t index = poc_get_index(indices, requested_index); | |
ChunkLocation resolved_index = index_resolver.Resolve(index); | |
int64_t chunk_index = resolved_index.chunk_index; | |
if (chunk_index < 0) { | |
// ChunkResolver doesn't throw errors when the index is out of bounds | |
// it will just return a chunk index that doesn't exist. | |
return Status::IndexError("Index ", index, " is out of bounds"); | |
} | |
ARROW_RETURN_NOT_OK(builders[chunk_index].Append(resolved_index.index_in_chunk)); | |
} | |
std::vector<std::shared_ptr<arrow::Array>> looked_up_values; | |
looked_up_values.reserve(num_chunks); | |
for (int i = 0; i < num_chunks; ++i) { | |
if (builders[i].length() == 0) { | |
// No indices refer to this chunk, so we can skip it | |
continue; | |
} | |
std::shared_ptr<Int64Array> indices_array; | |
ARROW_RETURN_NOT_OK(builders[i].Finish(&indices_array)); | |
std::shared_ptr<ArrayData> looked_up_values_data; | |
ARROW_ASSIGN_OR_RAISE(looked_up_values_data, TakeAA(values.chunk(i)->data(), | |
indices_array->data(), options, ctx)); | |
looked_up_values[i] = MakeArray(looked_up_values_data); | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment