You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@arrow.apache.org by "amol- (via GitHub)" <gi...@apache.org> on 2024/01/11 14:57:38 UTC
[PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
amol- opened a new pull request, #39566:
URL: https://github.com/apache/arrow/pull/39566
### Rationale for this change
We can avoid extra unecessary work and memory consumption of concatenating chunks when running take, we can directly run the take on the chunks at the only cost of remapping the indices which are usually much fewer than the size of the array we are applying take on.
### Are these changes tested?
Two tests already existed that verify take on ChunkedArray and they covered the corner cases well, the only tweak necessary to those tests was that now take returns a chunkedarray made of multiple chunks instead of a single one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "amol- (via GitHub)" <gi...@apache.org>.
amol- commented on PR #39566:
URL: https://github.com/apache/arrow/pull/39566#issuecomment-1969353021
> @felipecrv @amol- Should this PR be kept open now that #40206 was merged?
I think so, this PR is focused on optimizing TakeCA, while the one that was merged was focused on TakeCC
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "pitrou (via GitHub)" <gi...@apache.org>.
pitrou commented on PR #39566:
URL: https://github.com/apache/arrow/pull/39566#issuecomment-1969269062
@felipecrv @amol- Should this PR be kept open now that #40206 was merged?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "amol- (via GitHub)" <gi...@apache.org>.
amol- closed pull request #39566: GH-39565: [C++] Do not concatenate ChunkedArray when running take function
URL: https://github.com/apache/arrow/pull/39566
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "amol- (via GitHub)" <gi...@apache.org>.
amol- commented on code in PR #39566:
URL: https://github.com/apache/arrow/pull/39566#discussion_r1452392211
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -660,29 +660,95 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
auto num_chunks = values.num_chunks();
std::shared_ptr<Array> current_chunk;
- // Case 1: `values` has a single chunk, so just use it
- if (num_chunks == 1) {
- current_chunk = values.chunk(0);
- } else {
- // TODO Case 2: See if all `indices` fall in the same chunk and call Array Take on it
- // See
- // https://github.com/apache/arrow/blob/6f2c9041137001f7a9212f244b51bc004efc29af/r/src/compute.cpp#L123-L151
- // TODO Case 3: If indices are sorted, can slice them and call Array Take
-
- // Case 4: Else, concatenate chunks and call Array Take
+ if (indices.length() == 0) {
+ // Case 0: No indices were provided, nothing to take so return an empty chunked array
+ return ChunkedArray::MakeEmpty(values.type());
+ } else if (num_chunks < 2) {
+ // Case 1: `values` is empty or has a single chunk, so just use it
if (values.chunks().empty()) {
ARROW_ASSIGN_OR_RAISE(current_chunk, MakeArrayOfNull(values.type(), /*length=*/0,
ctx->memory_pool()));
} else {
- ARROW_ASSIGN_OR_RAISE(current_chunk,
- Concatenate(values.chunks(), ctx->memory_pool()));
+ current_chunk = values.chunk(0);
}
+ // Call Array Take on our single chunk
+ ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> new_chunk,
+ TakeAA(current_chunk->data(), indices.data(), options, ctx));
+ std::vector<std::shared_ptr<Array>> chunks = {MakeArray(new_chunk)};
+ return std::make_shared<ChunkedArray>(std::move(chunks));
+ } else {
+ // For each index, lookup at which chunk it refers to.
+ // We have to do this because the indices are not necessarily sorted.
+ // So we can't simply iterate over chunks and pick the slices we need.
+ std::vector<int64_t> indices_chunks(indices.length());
+ std::vector<int64_t> offsetted_indices(indices.length());
+ for (int64_t requested_index = 0; requested_index < indices.length();
+ ++requested_index) {
+ uint64_t index;
+ switch (indices.type()->byte_width()) {
Review Comment:
Migrated to `type()->id()`. Note that I'm still relying on the fact that the user will never pass negative indices as otherwise I would have to double the switch cases unless there is some kind of helper that can always return a uint64 for any provided integer value and error for values <0
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #39566:
URL: https://github.com/apache/arrow/pull/39566#discussion_r1467166688
##########
cpp/src/arrow/chunk_resolver.h:
##########
@@ -57,9 +60,14 @@ struct ARROW_EXPORT ChunkResolver {
// This is trivial when merging (assuming each side of the merge uses
// its own resolver), but also in the inner recursive invocations of
// partitioning.
+ if (index < 0 || index >= max_index_) {
+ return {-1, -1};
+ }
Review Comment:
Returning `offsets_.size()` as the `chunk_index` to indicate the index was not found is better:
1) `std::upper_bound`-like implementation can handle it naturally without an extra branch
2) the `.end()` iterator is a common way of representing that something was not found. Iterators, like indices represent the `[0, length)` range -- `begin = 0, end = length`
3) the negative value could be cast to unsigned and create annoying overflows
```diff
diff --git a/cpp/src/arrow/chunk_resolver.cc b/cpp/src/arrow/chunk_resolver.cc
index 45084f8f5..4a1ba6d0a 100644
--- a/cpp/src/arrow/chunk_resolver.cc
+++ b/cpp/src/arrow/chunk_resolver.cc
@@ -57,19 +57,13 @@ inline std::vector<int64_t> MakeChunksOffsets(const std::vector<T>& chunks) {
} // namespace
ChunkResolver::ChunkResolver(const ArrayVector& chunks)
- : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
ChunkResolver::ChunkResolver(const std::vector<const Array*>& chunks)
- : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
ChunkResolver::ChunkResolver(const RecordBatchVector& batches)
- : offsets_(MakeChunksOffsets(batches)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(batches)), cached_chunk_(0) {}
} // namespace internal
} // namespace arrow
diff --git a/cpp/src/arrow/chunk_resolver.h b/cpp/src/arrow/chunk_resolver.h
index aab63094d..ff0a0b437 100644
--- a/cpp/src/arrow/chunk_resolver.h
+++ b/cpp/src/arrow/chunk_resolver.h
@@ -40,19 +40,18 @@ struct ARROW_EXPORT ChunkResolver {
explicit ChunkResolver(const RecordBatchVector& batches);
ChunkResolver(ChunkResolver&& other) noexcept
- : offsets_(std::move(other.offsets_)),
- max_index_(other.max_index_),
- cached_chunk_(other.cached_chunk_.load()) {}
+ : offsets_(std::move(other.offsets_)), cached_chunk_(other.cached_chunk_.load()) {}
ChunkResolver& operator=(ChunkResolver&& other) {
offsets_ = std::move(other.offsets_);
- max_index_ = other.max_index_;
cached_chunk_.store(other.cached_chunk_.load());
return *this;
}
/// \brief Return a ChunkLocation containing the chunk index and in-chunk value index of
/// the chunked array at logical index
+ ///
+ /// \pre index >= 0
inline ChunkLocation Resolve(const int64_t index) const {
// It is common for the algorithms below to make consecutive accesses at
// a relatively small distance from each other, hence often falling in
@@ -60,14 +59,9 @@ struct ARROW_EXPORT ChunkResolver {
// This is trivial when merging (assuming each side of the merge uses
// its own resolver), but also in the inner recursive invocations of
// partitioning.
- if (index < 0 || index >= max_index_) {
- return {-1, -1};
- }
-
if (offsets_.size() <= 1) {
return {0, index};
}
-
const auto cached_chunk = cached_chunk_.load();
const bool cache_hit =
(index >= offsets_[cached_chunk] && index < offsets_[cached_chunk + 1]);
@@ -103,9 +97,6 @@ struct ARROW_EXPORT ChunkResolver {
// Collection of starting offsets used for binary search
std::vector<int64_t> offsets_;
- // Latest valid index
- int64_t max_index_;
-
// Tracks the most recently used chunk index to allow fast
// access for consecutive indices corresponding to the same chunk
mutable std::atomic<int64_t> cached_chunk_;
diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc
index 6a2911b90..c36b736d5 100644
--- a/cpp/src/arrow/chunked_array.cc
+++ b/cpp/src/arrow/chunked_array.cc
@@ -169,7 +169,7 @@ bool ChunkedArray::ApproxEquals(const ChunkedArray& other,
Result<std::shared_ptr<Scalar>> ChunkedArray::GetScalar(int64_t index) const {
const auto loc = chunk_resolver_.Resolve(index);
- if (loc.chunk_index == -1) {
+ if (loc.chunk_index >= static_cast<int64_t>(chunks_.size())) {
return Status::IndexError("index with value of ", index,
" is out-of-bounds for chunked array of length ", length_);
}
diff --git a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
index 3eefab610..f8e4d78fe 100644
--- a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
+++ b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
@@ -738,9 +738,7 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
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.
+ if (chunk_index >= num_chunks) {
return Status::IndexError("Index ", index, " is out of bounds");
}
indices_chunks[requested_index] = chunk_index;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39566:
URL: https://github.com/apache/arrow/pull/39566#issuecomment-1887360568
:warning: GitHub issue #39565 **has been automatically assigned in GitHub** to PR creator.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "benibus (via GitHub)" <gi...@apache.org>.
benibus commented on code in PR #39566:
URL: https://github.com/apache/arrow/pull/39566#discussion_r1449197127
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -660,29 +660,95 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
auto num_chunks = values.num_chunks();
std::shared_ptr<Array> current_chunk;
- // Case 1: `values` has a single chunk, so just use it
- if (num_chunks == 1) {
- current_chunk = values.chunk(0);
- } else {
- // TODO Case 2: See if all `indices` fall in the same chunk and call Array Take on it
- // See
- // https://github.com/apache/arrow/blob/6f2c9041137001f7a9212f244b51bc004efc29af/r/src/compute.cpp#L123-L151
- // TODO Case 3: If indices are sorted, can slice them and call Array Take
-
- // Case 4: Else, concatenate chunks and call Array Take
+ if (indices.length() == 0) {
+ // Case 0: No indices were provided, nothing to take so return an empty chunked array
+ return ChunkedArray::MakeEmpty(values.type());
+ } else if (num_chunks < 2) {
+ // Case 1: `values` is empty or has a single chunk, so just use it
if (values.chunks().empty()) {
ARROW_ASSIGN_OR_RAISE(current_chunk, MakeArrayOfNull(values.type(), /*length=*/0,
ctx->memory_pool()));
} else {
- ARROW_ASSIGN_OR_RAISE(current_chunk,
- Concatenate(values.chunks(), ctx->memory_pool()));
+ current_chunk = values.chunk(0);
}
+ // Call Array Take on our single chunk
+ ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> new_chunk,
+ TakeAA(current_chunk->data(), indices.data(), options, ctx));
+ std::vector<std::shared_ptr<Array>> chunks = {MakeArray(new_chunk)};
+ return std::make_shared<ChunkedArray>(std::move(chunks));
+ } else {
+ // For each index, lookup at which chunk it refers to.
Review Comment:
For this routine, you should be able to use the [`ChunkResolver`](https://github.com/amol-/arrow/blob/1c2a2c833811aadf8c6887f807d607fc59f5c9b0/cpp/src/arrow/chunk_resolver.h#L35) utility to convert from logical indices to chunk/index pairs. As a bonus, it will cache the last used chunk, which could optimize sorted cases.
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -660,29 +660,95 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
auto num_chunks = values.num_chunks();
std::shared_ptr<Array> current_chunk;
- // Case 1: `values` has a single chunk, so just use it
- if (num_chunks == 1) {
- current_chunk = values.chunk(0);
- } else {
- // TODO Case 2: See if all `indices` fall in the same chunk and call Array Take on it
- // See
- // https://github.com/apache/arrow/blob/6f2c9041137001f7a9212f244b51bc004efc29af/r/src/compute.cpp#L123-L151
- // TODO Case 3: If indices are sorted, can slice them and call Array Take
-
- // Case 4: Else, concatenate chunks and call Array Take
+ if (indices.length() == 0) {
+ // Case 0: No indices were provided, nothing to take so return an empty chunked array
+ return ChunkedArray::MakeEmpty(values.type());
+ } else if (num_chunks < 2) {
+ // Case 1: `values` is empty or has a single chunk, so just use it
if (values.chunks().empty()) {
ARROW_ASSIGN_OR_RAISE(current_chunk, MakeArrayOfNull(values.type(), /*length=*/0,
ctx->memory_pool()));
} else {
- ARROW_ASSIGN_OR_RAISE(current_chunk,
- Concatenate(values.chunks(), ctx->memory_pool()));
+ current_chunk = values.chunk(0);
}
+ // Call Array Take on our single chunk
+ ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> new_chunk,
+ TakeAA(current_chunk->data(), indices.data(), options, ctx));
+ std::vector<std::shared_ptr<Array>> chunks = {MakeArray(new_chunk)};
+ return std::make_shared<ChunkedArray>(std::move(chunks));
+ } else {
+ // For each index, lookup at which chunk it refers to.
+ // We have to do this because the indices are not necessarily sorted.
+ // So we can't simply iterate over chunks and pick the slices we need.
+ std::vector<int64_t> indices_chunks(indices.length());
+ std::vector<int64_t> offsetted_indices(indices.length());
+ for (int64_t requested_index = 0; requested_index < indices.length();
+ ++requested_index) {
+ uint64_t index;
+ switch (indices.type()->byte_width()) {
Review Comment:
Seems a bit strange to switch on the byte width rather than `indices.type()->id()`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "amol- (via GitHub)" <gi...@apache.org>.
amol- commented on code in PR #39566:
URL: https://github.com/apache/arrow/pull/39566#discussion_r1452369504
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -660,29 +660,95 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
auto num_chunks = values.num_chunks();
std::shared_ptr<Array> current_chunk;
- // Case 1: `values` has a single chunk, so just use it
- if (num_chunks == 1) {
- current_chunk = values.chunk(0);
- } else {
- // TODO Case 2: See if all `indices` fall in the same chunk and call Array Take on it
- // See
- // https://github.com/apache/arrow/blob/6f2c9041137001f7a9212f244b51bc004efc29af/r/src/compute.cpp#L123-L151
- // TODO Case 3: If indices are sorted, can slice them and call Array Take
-
- // Case 4: Else, concatenate chunks and call Array Take
+ if (indices.length() == 0) {
+ // Case 0: No indices were provided, nothing to take so return an empty chunked array
+ return ChunkedArray::MakeEmpty(values.type());
+ } else if (num_chunks < 2) {
+ // Case 1: `values` is empty or has a single chunk, so just use it
if (values.chunks().empty()) {
ARROW_ASSIGN_OR_RAISE(current_chunk, MakeArrayOfNull(values.type(), /*length=*/0,
ctx->memory_pool()));
} else {
- ARROW_ASSIGN_OR_RAISE(current_chunk,
- Concatenate(values.chunks(), ctx->memory_pool()));
+ current_chunk = values.chunk(0);
}
+ // Call Array Take on our single chunk
+ ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ArrayData> new_chunk,
+ TakeAA(current_chunk->data(), indices.data(), options, ctx));
+ std::vector<std::shared_ptr<Array>> chunks = {MakeArray(new_chunk)};
+ return std::make_shared<ChunkedArray>(std::move(chunks));
+ } else {
+ // For each index, lookup at which chunk it refers to.
Review Comment:
Amazing, I didn't know about `ChunkResolver` seems super useful.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on PR #39566:
URL: https://github.com/apache/arrow/pull/39566#issuecomment-1969434210
Before my PR: `TakeCC` made `num_chunks` `Concatenate(chunks)` calls.
After my PR: `TakeCC` makes `1` `Concatenate(chunks)` call.
Next step (and goal of amol's PR/issue pair): 0 concatenations.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on PR #39566:
URL: https://github.com/apache/arrow/pull/39566#issuecomment-2116584150
I opened #41700 which can handle all the fixed-width types without concatenation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] GH-39565: [C++] Do not concatenate ChunkedArray when running take function [arrow]
Posted by "felipecrv (via GitHub)" <gi...@apache.org>.
felipecrv commented on code in PR #39566:
URL: https://github.com/apache/arrow/pull/39566#discussion_r1467166688
##########
cpp/src/arrow/chunk_resolver.h:
##########
@@ -57,9 +60,14 @@ struct ARROW_EXPORT ChunkResolver {
// This is trivial when merging (assuming each side of the merge uses
// its own resolver), but also in the inner recursive invocations of
// partitioning.
+ if (index < 0 || index >= max_index_) {
+ return {-1, -1};
+ }
Review Comment:
Returning `offsets_.size()` as the `chunk_index` to indicate the index was not found is better:
1) `upper_bound` implementation can handle it naturally without an extra branch
2) the `.end()` iterator is a common way of representing that something was not found. Iterators, like indices represent the `[0, length)` range -- `begin = 0, end = length`
3) the negative value could be cast to unsigned and create annoying overflows
```diff
diff --git a/cpp/src/arrow/chunk_resolver.cc b/cpp/src/arrow/chunk_resolver.cc
index 45084f8f5..4a1ba6d0a 100644
--- a/cpp/src/arrow/chunk_resolver.cc
+++ b/cpp/src/arrow/chunk_resolver.cc
@@ -57,19 +57,13 @@ inline std::vector<int64_t> MakeChunksOffsets(const std::vector<T>& chunks) {
} // namespace
ChunkResolver::ChunkResolver(const ArrayVector& chunks)
- : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
ChunkResolver::ChunkResolver(const std::vector<const Array*>& chunks)
- : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(chunks)), cached_chunk_(0) {}
ChunkResolver::ChunkResolver(const RecordBatchVector& batches)
- : offsets_(MakeChunksOffsets(batches)), cached_chunk_(0) {
- max_index_ = offsets_.back();
-}
+ : offsets_(MakeChunksOffsets(batches)), cached_chunk_(0) {}
} // namespace internal
} // namespace arrow
diff --git a/cpp/src/arrow/chunk_resolver.h b/cpp/src/arrow/chunk_resolver.h
index aab63094d..ff0a0b437 100644
--- a/cpp/src/arrow/chunk_resolver.h
+++ b/cpp/src/arrow/chunk_resolver.h
@@ -40,19 +40,18 @@ struct ARROW_EXPORT ChunkResolver {
explicit ChunkResolver(const RecordBatchVector& batches);
ChunkResolver(ChunkResolver&& other) noexcept
- : offsets_(std::move(other.offsets_)),
- max_index_(other.max_index_),
- cached_chunk_(other.cached_chunk_.load()) {}
+ : offsets_(std::move(other.offsets_)), cached_chunk_(other.cached_chunk_.load()) {}
ChunkResolver& operator=(ChunkResolver&& other) {
offsets_ = std::move(other.offsets_);
- max_index_ = other.max_index_;
cached_chunk_.store(other.cached_chunk_.load());
return *this;
}
/// \brief Return a ChunkLocation containing the chunk index and in-chunk value index of
/// the chunked array at logical index
+ ///
+ /// \pre index >= 0
inline ChunkLocation Resolve(const int64_t index) const {
// It is common for the algorithms below to make consecutive accesses at
// a relatively small distance from each other, hence often falling in
@@ -60,14 +59,9 @@ struct ARROW_EXPORT ChunkResolver {
// This is trivial when merging (assuming each side of the merge uses
// its own resolver), but also in the inner recursive invocations of
// partitioning.
- if (index < 0 || index >= max_index_) {
- return {-1, -1};
- }
-
if (offsets_.size() <= 1) {
return {0, index};
}
-
const auto cached_chunk = cached_chunk_.load();
const bool cache_hit =
(index >= offsets_[cached_chunk] && index < offsets_[cached_chunk + 1]);
@@ -103,9 +97,6 @@ struct ARROW_EXPORT ChunkResolver {
// Collection of starting offsets used for binary search
std::vector<int64_t> offsets_;
- // Latest valid index
- int64_t max_index_;
-
// Tracks the most recently used chunk index to allow fast
// access for consecutive indices corresponding to the same chunk
mutable std::atomic<int64_t> cached_chunk_;
diff --git a/cpp/src/arrow/chunked_array.cc b/cpp/src/arrow/chunked_array.cc
index 6a2911b90..c36b736d5 100644
--- a/cpp/src/arrow/chunked_array.cc
+++ b/cpp/src/arrow/chunked_array.cc
@@ -169,7 +169,7 @@ bool ChunkedArray::ApproxEquals(const ChunkedArray& other,
Result<std::shared_ptr<Scalar>> ChunkedArray::GetScalar(int64_t index) const {
const auto loc = chunk_resolver_.Resolve(index);
- if (loc.chunk_index == -1) {
+ if (loc.chunk_index >= static_cast<int64_t>(chunks_.size())) {
return Status::IndexError("index with value of ", index,
" is out-of-bounds for chunked array of length ", length_);
}
diff --git a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
index 3eefab610..f8e4d78fe 100644
--- a/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
+++ b/cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc
@@ -738,9 +738,7 @@ Result<std::shared_ptr<ChunkedArray>> TakeCA(const ChunkedArray& values,
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.
+ if (chunk_index >= num_chunks) {
return Status::IndexError("Index ", index, " is out of bounds");
}
indices_chunks[requested_index] = chunk_index;
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org