GH-48277: [C++][Parquet] unpack with shuffle algorithm

cpp/src/arrow/util/bit_stream_utils_internal.h

@@ -203,8 +203,9 @@ inline bool BitWriter::PutValue(uint64_t v, int num_bits) {
 
   if (ARROW_PREDICT_FALSE(static_cast<int64_t>(byte_offset_) * 8 + bit_offset_ +
                               num_bits >
-                          static_cast<int64_t>(max_bytes_) * 8))
+                          static_cast<int64_t>(max_bytes_) * 8)) {
     return false;
+  }
 
   buffered_values_ |= v << bit_offset_;
   bit_offset_ += num_bits;
@@ -273,14 +274,19 @@ inline int BitReader::GetBatch(int num_bits, T* v, int batch_size) {
     batch_size = static_cast<int>(remaining_bits / num_bits);
   }
 
+  const ::arrow::internal::UnpackOptions opts{
+      .batch_size = batch_size,
+      .bit_width = num_bits,
+      .bit_offset = bit_offset_,
+      .max_read_bytes = max_bytes_ - byte_offset_,
+  };
+
   if constexpr (std::is_same_v<T, bool>) {
-    ::arrow::internal::unpack(buffer_ + byte_offset_, v, batch_size, num_bits,
-                              bit_offset_);
+    ::arrow::internal::unpack(buffer_ + byte_offset_, v, opts);
 
   } else {
     ::arrow::internal::unpack(buffer_ + byte_offset_,
-                              reinterpret_cast<std::make_unsigned_t<T>*>(v), batch_size,
-                              num_bits, bit_offset_);
+                              reinterpret_cast<std::make_unsigned_t<T>*>(v), opts);
   }
 
   Advance(batch_size * num_bits);

cpp/src/arrow/util/bpacking.cc

@@ -17,7 +17,6 @@
 
 #include <array>
 
-#include "arrow/util/bpacking_dispatch_internal.h"
 #include "arrow/util/bpacking_internal.h"
 #include "arrow/util/bpacking_scalar_internal.h"
 #include "arrow/util/bpacking_simd_internal.h"
@@ -29,19 +28,21 @@ namespace {
 
 template <typename Uint>
 struct UnpackDynamicFunction {
-  using FunctionType = decltype(&unpack_scalar<Uint>);
+  using FunctionType = decltype(&bpacking::unpack_scalar<Uint>);
   using Implementation = std::pair<DispatchLevel, FunctionType>;
 
   static constexpr auto implementations() {
     return std::array{
-        // Current SIMD unpack algorithm works terribly on SSE4.2 due to lack of variable
-        // rhsift and poor xsimd fallback.
-        Implementation{DispatchLevel::NONE, &unpack_scalar<Uint>},
+#if defined(ARROW_HAVE_SSE4_2)
+        Implementation{DispatchLevel::NONE, &bpacking::unpack_sse4_2<Uint>},
+#else
+        Implementation{DispatchLevel::NONE, &bpacking::unpack_scalar<Uint>},
+#endif
 #if defined(ARROW_HAVE_RUNTIME_AVX2)
-        Implementation{DispatchLevel::AVX2, &unpack_avx2<Uint>},
+        Implementation{DispatchLevel::AVX2, &bpacking::unpack_avx2<Uint>},
 #endif
 #if defined(ARROW_HAVE_RUNTIME_AVX512)
-        Implementation{DispatchLevel::AVX512, &unpack_avx512<Uint>},
+        Implementation{DispatchLevel::AVX512, &bpacking::unpack_avx512<Uint>},
 #endif
     };
   }
@@ -50,19 +51,19 @@ struct UnpackDynamicFunction {
 }  // namespace
 
 template <typename Uint>
-void unpack(const uint8_t* in, Uint* out, int batch_size, int num_bits, int bit_offset) {
+void unpack(const uint8_t* in, Uint* out, const UnpackOptions& opts) {
 #if defined(ARROW_HAVE_NEON)
-  return unpack_neon(in, out, batch_size, num_bits, bit_offset);
+  return bpacking::unpack_neon(in, out, opts);
 #else
   static DynamicDispatch<UnpackDynamicFunction<Uint> > dispatch;
-  return dispatch.func(in, out, batch_size, num_bits, bit_offset);
+  return dispatch.func(in, out, opts);
 #endif
 }
 
-template void unpack<bool>(const uint8_t*, bool*, int, int, int);
-template void unpack<uint8_t>(const uint8_t*, uint8_t*, int, int, int);
-template void unpack<uint16_t>(const uint8_t*, uint16_t*, int, int, int);
-template void unpack<uint32_t>(const uint8_t*, uint32_t*, int, int, int);
-template void unpack<uint64_t>(const uint8_t*, uint64_t*, int, int, int);
+template void unpack<bool>(const uint8_t*, bool*, const UnpackOptions&);
+template void unpack<uint8_t>(const uint8_t*, uint8_t*, const UnpackOptions&);
+template void unpack<uint16_t>(const uint8_t*, uint16_t*, const UnpackOptions&);
+template void unpack<uint32_t>(const uint8_t*, uint32_t*, const UnpackOptions&);
+template void unpack<uint64_t>(const uint8_t*, uint64_t*, const UnpackOptions&);
 
 }  // namespace arrow::internal

cpp/src/arrow/util/bpacking_benchmark.cc

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+#include <memory>
 #include <stdexcept>
 #include <vector>
 
@@ -33,7 +34,7 @@ namespace arrow::internal {
 namespace {
 
 template <typename Int>
-using UnpackFunc = void (*)(const uint8_t*, Int*, int, int, int);
+using UnpackFunc = void (*)(const uint8_t*, Int*, const UnpackOptions&);
 
 /// Get the number of bytes associate with a packing.
 constexpr int32_t GetNumBytes(int32_t num_values, int32_t bit_width) {
@@ -86,33 +87,62 @@ void BM_Unpack(benchmark::State& state, bool aligned, UnpackFunc<Int> unpack, bo
   const uint8_t* packed_ptr =
       GetNextAlignedByte(packed.data(), sizeof(Int)) + (aligned ? 0 : 1);
 
-  std::vector<Int> unpacked(num_values, 0);
+  auto unpacked = std::make_unique<Int[]>(num_values);
+
+  const ::arrow::internal::UnpackOptions opts{
+      .batch_size = num_values,
+      .bit_width = bit_width,
+      .bit_offset = 0,
+      .max_read_bytes = -1,
+  };
 
   for (auto _ : state) {
-    unpack(packed_ptr, unpacked.data(), num_values, bit_width, /* bit_offset = */ 0);
+    unpack(packed_ptr, unpacked.get(), opts);
     benchmark::ClobberMemory();
   }
   state.SetItemsProcessed(num_values * state.iterations());
 }
 
-constexpr int32_t kMinRange = 64;
-constexpr int32_t kMaxRange = 32768;
+// Currently, the minimum unpack SIMD kernel size is 32 and the RLE-bit-packing encoder
+// will not emit runs larger than 512 (though other implementation might), so we biased
+// the benchmarks towards a rather small scale.
+static const auto kNumValuesRange = benchmark::CreateRange(32, 512, 2);
+constexpr std::initializer_list<int64_t> kBitWidths8 = {1, 2, 8};
 constexpr std::initializer_list<int64_t> kBitWidths16 = {1, 2, 8, 13};
 constexpr std::initializer_list<int64_t> kBitWidths32 = {1, 2, 8, 20};
 constexpr std::initializer_list<int64_t> kBitWidths64 = {1, 2, 8, 20, 47};
+
+static const std::vector<std::vector<int64_t>> kBitWidthsNumValuesBool = {
+    {0, 1},
+    kNumValuesRange,
+};
+static const std::vector<std::vector<int64_t>> kBitWidthsNumValues8 = {
+    kBitWidths8,
+    kNumValuesRange,
+};
 static const std::vector<std::vector<int64_t>> kBitWidthsNumValues16 = {
     kBitWidths16,
-    benchmark::CreateRange(kMinRange, kMaxRange, /*multi=*/32),
+    kNumValuesRange,
 };
 static const std::vector<std::vector<int64_t>> kBitWidthsNumValues32 = {
     kBitWidths32,
-    benchmark::CreateRange(kMinRange, kMaxRange, /*multi=*/32),
+    kNumValuesRange,
 };
 static const std::vector<std::vector<int64_t>> kBitWidthsNumValues64 = {
     kBitWidths64,
-    benchmark::CreateRange(kMinRange, kMaxRange, /*multi=*/32),
+    kNumValuesRange,
 };
 
+/// Nudge for MSVC template inside BENCHMARK_CAPTURE macro.
+void BM_UnpackBool(benchmark::State& state, bool aligned, UnpackFunc<bool> unpack,
+                   bool skip = false, std::string skip_msg = "") {
+  return BM_Unpack<bool>(state, aligned, unpack, skip, std::move(skip_msg));
+}
+/// Nudge for MSVC template inside BENCHMARK_CAPTURE macro.
+void BM_UnpackUint8(benchmark::State& state, bool aligned, UnpackFunc<uint8_t> unpack,
+                    bool skip = false, std::string skip_msg = "") {
+  return BM_Unpack<uint8_t>(state, aligned, unpack, skip, std::move(skip_msg));
+}
 /// Nudge for MSVC template inside BENCHMARK_CAPTURE macro.
 void BM_UnpackUint16(benchmark::State& state, bool aligned, UnpackFunc<uint16_t> unpack,
                      bool skip = false, std::string skip_msg = "") {
@@ -129,52 +159,111 @@ void BM_UnpackUint64(benchmark::State& state, bool aligned, UnpackFunc<uint64_t>
   return BM_Unpack<uint64_t>(state, aligned, unpack, skip, std::move(skip_msg));
 }
 
-BENCHMARK_CAPTURE(BM_UnpackUint16, ScalarUnaligned, false, &unpack_scalar<uint16_t>)
+BENCHMARK_CAPTURE(BM_UnpackBool, ScalarUnaligned, false, &bpacking::unpack_scalar<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, ScalarUnaligned, false,
+                  &bpacking::unpack_scalar<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, ScalarUnaligned, false,
+                  &bpacking::unpack_scalar<uint16_t>)
     ->ArgsProduct(kBitWidthsNumValues16);
-BENCHMARK_CAPTURE(BM_UnpackUint32, ScalarUnaligned, false, &unpack_scalar<uint32_t>)
+BENCHMARK_CAPTURE(BM_UnpackUint32, ScalarUnaligned, false,
+                  &bpacking::unpack_scalar<uint32_t>)
     ->ArgsProduct(kBitWidthsNumValues32);
-BENCHMARK_CAPTURE(BM_UnpackUint64, ScalarUnaligned, false, &unpack_scalar<uint64_t>)
+BENCHMARK_CAPTURE(BM_UnpackUint64, ScalarUnaligned, false,
+                  &bpacking::unpack_scalar<uint64_t>)
     ->ArgsProduct(kBitWidthsNumValues64);
 
+#if defined(ARROW_HAVE_SSE4_2)
+BENCHMARK_CAPTURE(BM_UnpackBool, Sse42Unaligned, false, &bpacking::unpack_sse4_2<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Sse42Unaligned, false,
+                  &bpacking::unpack_sse4_2<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, Sse42Unaligned, false,
+                  &bpacking::unpack_sse4_2<uint16_t>)
+    ->ArgsProduct(kBitWidthsNumValues16);
+BENCHMARK_CAPTURE(BM_UnpackUint32, Sse42Unaligned, false,
+                  &bpacking::unpack_sse4_2<uint32_t>)
+    ->ArgsProduct(kBitWidthsNumValues32);
+BENCHMARK_CAPTURE(BM_UnpackUint64, Sse42Unaligned, false,
+                  &bpacking::unpack_sse4_2<uint64_t>)
+    ->ArgsProduct(kBitWidthsNumValues64);
+#endif
+
 #if defined(ARROW_HAVE_RUNTIME_AVX2)
-BENCHMARK_CAPTURE(BM_UnpackUint16, Avx2Unaligned, false, &unpack_avx2<uint16_t>,
+BENCHMARK_CAPTURE(BM_UnpackBool, Avx2Unaligned, false, &bpacking::unpack_avx2<bool>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
+                  "Avx2 not available")
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Avx2Unaligned, false, &bpacking::unpack_avx2<uint8_t>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
+                  "Avx2 not available")
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, Avx2Unaligned, false, &bpacking::unpack_avx2<uint16_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
                   "Avx2 not available")
     ->ArgsProduct(kBitWidthsNumValues16);
-BENCHMARK_CAPTURE(BM_UnpackUint32, Avx2Unaligned, false, &unpack_avx2<uint32_t>,
+BENCHMARK_CAPTURE(BM_UnpackUint32, Avx2Unaligned, false, &bpacking::unpack_avx2<uint32_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
                   "Avx2 not available")
     ->ArgsProduct(kBitWidthsNumValues32);
-BENCHMARK_CAPTURE(BM_UnpackUint64, Avx2Unaligned, false, &unpack_avx2<uint64_t>,
+BENCHMARK_CAPTURE(BM_UnpackUint64, Avx2Unaligned, false, &bpacking::unpack_avx2<uint64_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
                   "Avx2 not available")
     ->ArgsProduct(kBitWidthsNumValues64);
 #endif
 
 #if defined(ARROW_HAVE_RUNTIME_AVX512)
-BENCHMARK_CAPTURE(BM_UnpackUint16, Avx512Unaligned, false, &unpack_avx512<uint16_t>,
+BENCHMARK_CAPTURE(BM_UnpackBool, Avx512Unaligned, false, &bpacking::unpack_avx512<bool>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
+                  "Avx512 not available")
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Avx512Unaligned, false,
+                  &bpacking::unpack_avx512<uint8_t>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
+                  "Avx512 not available")
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, Avx512Unaligned, false,
+                  &bpacking::unpack_avx512<uint16_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
                   "Avx512 not available")
     ->ArgsProduct(kBitWidthsNumValues16);
-BENCHMARK_CAPTURE(BM_UnpackUint32, Avx512Unaligned, false, &unpack_avx512<uint32_t>,
+BENCHMARK_CAPTURE(BM_UnpackUint32, Avx512Unaligned, false,
+                  &bpacking::unpack_avx512<uint32_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
                   "Avx512 not available")
     ->ArgsProduct(kBitWidthsNumValues32);
-BENCHMARK_CAPTURE(BM_UnpackUint64, Avx512Unaligned, false, &unpack_avx512<uint64_t>,
+BENCHMARK_CAPTURE(BM_UnpackUint64, Avx512Unaligned, false,
+                  &bpacking::unpack_avx512<uint64_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
                   "Avx512 not available")
     ->ArgsProduct(kBitWidthsNumValues64);
 #endif
 
 #if defined(ARROW_HAVE_NEON)
-BENCHMARK_CAPTURE(BM_UnpackUint16, NeonUnaligned, false, &unpack_neon<uint16_t>)
+BENCHMARK_CAPTURE(BM_UnpackBool, NeonUnaligned, false, &bpacking::unpack_neon<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, NeonUnaligned, false, &bpacking::unpack_neon<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, NeonUnaligned, false, &bpacking::unpack_neon<uint16_t>)
     ->ArgsProduct(kBitWidthsNumValues16);
-BENCHMARK_CAPTURE(BM_UnpackUint32, NeonUnaligned, false, &unpack_neon<uint32_t>)
+BENCHMARK_CAPTURE(BM_UnpackUint32, NeonUnaligned, false, &bpacking::unpack_neon<uint32_t>)
     ->ArgsProduct(kBitWidthsNumValues32);
-BENCHMARK_CAPTURE(BM_UnpackUint64, NeonUnaligned, false, &unpack_neon<uint64_t>)
+BENCHMARK_CAPTURE(BM_UnpackUint64, NeonUnaligned, false, &bpacking::unpack_neon<uint64_t>)
     ->ArgsProduct(kBitWidthsNumValues64);
 #endif
 
+BENCHMARK_CAPTURE(BM_UnpackBool, DynamicAligned, true, &unpack<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackBool, DynamicUnaligned, false, &unpack<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+
+BENCHMARK_CAPTURE(BM_UnpackUint8, DynamicAligned, true, &unpack<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint8, DynamicUnaligned, false, &unpack<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+
 BENCHMARK_CAPTURE(BM_UnpackUint16, DynamicAligned, true, &unpack<uint16_t>)
     ->ArgsProduct(kBitWidthsNumValues16);
 BENCHMARK_CAPTURE(BM_UnpackUint16, DynamicUnaligned, false, &unpack<uint16_t>)