felipecrv commented on code in PR #41373:
URL: https://github.com/apache/arrow/pull/41373#discussion_r1591380457
##########
cpp/src/arrow/compute/kernels/vector_selection_take_internal.cc:
##########
@@ -324,261 +326,109 @@ namespace {
using TakeState = OptionsWrapper<TakeOptions>;
// ----------------------------------------------------------------------
-// Implement optimized take for primitive types from boolean to
1/2/4/8/16/32-byte
-// C-type based types. Use common implementation for every byte width and only
-// generate code for unsigned integer indices, since after boundschecking to
-// check for negative numbers in the indices we can safely reinterpret_cast
-// signed integers as unsigned.
-
-/// \brief The Take implementation for primitive (fixed-width) types does not
-/// use the logical Arrow type but rather the physical C type. This way we
-/// only generate one take function for each byte width.
+// Implement optimized take for primitive types from boolean to
+// 1/2/4/8/16/32-byte C-type based types and fixed-size binary (0 or more
+// bytes).
+//
+// Use one specialization for each of these primitive byte-widths so the
+// compiler can specialize the memcpy to dedicated CPU instructions and for
+// fixed-width binary use the 1-byte specialization but pass WithFactor=true
+// that makes the kernel consider the factor parameter provided at runtime.
+//
+// Only unsigned index types need to be instantiated since after
+// boundschecking to check for negative numbers in the indices we can safely
+// reinterpret_cast signed integers as unsigned.
+
+/// \brief The Take implementation for primitive types and fixed-width binary.
///
/// Also note that this function can also handle fixed-size-list arrays if
/// they fit the criteria described in fixed_width_internal.h, so use the
/// function defined in that file to access values and destination pointers
/// and DO NOT ASSUME `values.type()` is a primitive type.
///
/// \pre the indices have been boundschecked
-template <typename IndexCType, typename ValueWidthConstant>
-struct PrimitiveTakeImpl {
- static constexpr int kValueWidth = ValueWidthConstant::value;
-
- static void Exec(const ArraySpan& values, const ArraySpan& indices,
- ArrayData* out_arr) {
- DCHECK_EQ(util::FixedWidthInBytes(*values.type), kValueWidth);
- const auto* values_data = util::OffsetPointerOfFixedWidthValues(values);
- const uint8_t* values_is_valid = values.buffers[0].data;
- auto values_offset = values.offset;
-
- const auto* indices_data = indices.GetValues<IndexCType>(1);
- const uint8_t* indices_is_valid = indices.buffers[0].data;
- auto indices_offset = indices.offset;
-
- DCHECK_EQ(out_arr->offset, 0);
- auto* out = util::MutableFixedWidthValuesPointer(out_arr);
- auto out_is_valid = out_arr->buffers[0]->mutable_data();
-
- // If either the values or indices have nulls, we preemptively zero out the
- // out validity bitmap so that we don't have to use ClearBit in each
- // iteration for nulls.
- if (values.null_count != 0 || indices.null_count != 0) {
- bit_util::SetBitsTo(out_is_valid, 0, indices.length, false);
- }
-
- auto WriteValue = [&](int64_t position) {
- memcpy(out + position * kValueWidth,
- values_data + indices_data[position] * kValueWidth, kValueWidth);
- };
-
- auto WriteZero = [&](int64_t position) {
- memset(out + position * kValueWidth, 0, kValueWidth);
- };
-
- auto WriteZeroSegment = [&](int64_t position, int64_t length) {
- memset(out + position * kValueWidth, 0, kValueWidth * length);
- };
-
- OptionalBitBlockCounter indices_bit_counter(indices_is_valid,
indices_offset,
- indices.length);
- int64_t position = 0;
- int64_t valid_count = 0;
- while (position < indices.length) {
- BitBlockCount block = indices_bit_counter.NextBlock();
- if (values.null_count == 0) {
- // Values are never null, so things are easier
- valid_count += block.popcount;
- if (block.popcount == block.length) {
- // Fastest path: neither values nor index nulls
- bit_util::SetBitsTo(out_is_valid, position, block.length, true);
- for (int64_t i = 0; i < block.length; ++i) {
- WriteValue(position);
- ++position;
- }
- } else if (block.popcount > 0) {
- // Slow path: some indices but not all are null
- for (int64_t i = 0; i < block.length; ++i) {
- if (bit_util::GetBit(indices_is_valid, indices_offset + position))
{
- // index is not null
- bit_util::SetBit(out_is_valid, position);
- WriteValue(position);
- } else {
- WriteZero(position);
- }
- ++position;
- }
- } else {
- WriteZeroSegment(position, block.length);
- position += block.length;
- }
- } else {
- // Values have nulls, so we must do random access into the values
bitmap
- if (block.popcount == block.length) {
- // Faster path: indices are not null but values may be
- for (int64_t i = 0; i < block.length; ++i) {
- if (bit_util::GetBit(values_is_valid,
- values_offset + indices_data[position])) {
- // value is not null
- WriteValue(position);
- bit_util::SetBit(out_is_valid, position);
- ++valid_count;
- } else {
- WriteZero(position);
- }
- ++position;
- }
- } else if (block.popcount > 0) {
- // Slow path: some but not all indices are null. Since we are doing
- // random access in general we have to check the value nullness one
by
- // one.
- for (int64_t i = 0; i < block.length; ++i) {
- if (bit_util::GetBit(indices_is_valid, indices_offset + position)
&&
- bit_util::GetBit(values_is_valid,
- values_offset + indices_data[position])) {
- // index is not null && value is not null
- WriteValue(position);
- bit_util::SetBit(out_is_valid, position);
- ++valid_count;
- } else {
- WriteZero(position);
- }
- ++position;
- }
- } else {
- WriteZeroSegment(position, block.length);
- position += block.length;
- }
- }
+template <typename IndexCType, typename ValueBitWidthConstant,
+ typename OutputIsZeroInitialized = std::false_type,
+ typename WithFactor = std::false_type>
Review Comment:
No specific reason.
I will change, as I prefer to use `bool` directly myself. I was trying to be
compatible with the use of `std::integral_constant<>` for
`ValueBitWidthConstant` that was added by you in a previous PR :)
What's the reasoning for using `std::integral_constant` instead of an
integer typed template parameter?
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
