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

cpp/src/arrow/util/bit_stream_utils_internal.h

@@ -273,14 +273,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"
@@ -34,9 +33,11 @@ struct UnpackDynamicFunction {
 
   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.
+#if defined(ARROW_HAVE_SSE4_2)
+      Implementation{DispatchLevel::NONE, &unpack_sse4_2<Uint>},
+#else
       Implementation{DispatchLevel::NONE, &unpack_scalar<Uint>},
+#endif
 #if defined(ARROW_HAVE_RUNTIME_AVX2)
           Implementation{DispatchLevel::AVX2, &unpack_avx2<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 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 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,14 +159,39 @@ 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_UnpackBool, ScalarUnaligned, false, &unpack_scalar<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, ScalarUnaligned, false, &unpack_scalar<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
 BENCHMARK_CAPTURE(BM_UnpackUint16, ScalarUnaligned, false, &unpack_scalar<uint16_t>)
     ->ArgsProduct(kBitWidthsNumValues16);
 BENCHMARK_CAPTURE(BM_UnpackUint32, ScalarUnaligned, false, &unpack_scalar<uint32_t>)
     ->ArgsProduct(kBitWidthsNumValues32);
 BENCHMARK_CAPTURE(BM_UnpackUint64, ScalarUnaligned, false, &unpack_scalar<uint64_t>)
     ->ArgsProduct(kBitWidthsNumValues64);
 
+#if defined(ARROW_HAVE_SSE4_2)
+BENCHMARK_CAPTURE(BM_UnpackBool, Sse42Unaligned, false, &unpack_sse4_2<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Sse42Unaligned, false, &unpack_sse4_2<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
+BENCHMARK_CAPTURE(BM_UnpackUint16, Sse42Unaligned, false, &unpack_sse4_2<uint16_t>)
+    ->ArgsProduct(kBitWidthsNumValues16);
+BENCHMARK_CAPTURE(BM_UnpackUint32, Sse42Unaligned, false, &unpack_sse4_2<uint32_t>)
+    ->ArgsProduct(kBitWidthsNumValues32);
+BENCHMARK_CAPTURE(BM_UnpackUint64, Sse42Unaligned, false, &unpack_sse4_2<uint64_t>)
+    ->ArgsProduct(kBitWidthsNumValues64);
+#endif
+
 #if defined(ARROW_HAVE_RUNTIME_AVX2)
+BENCHMARK_CAPTURE(BM_UnpackBool, Avx2Unaligned, false, &unpack_avx2<bool>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
+                  "Avx2 not available")
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Avx2Unaligned, false, &unpack_avx2<uint8_t>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
+                  "Avx2 not available")
+    ->ArgsProduct(kBitWidthsNumValues8);
 BENCHMARK_CAPTURE(BM_UnpackUint16, Avx2Unaligned, false, &unpack_avx2<uint16_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX2),
                   "Avx2 not available")
@@ -152,6 +207,14 @@ BENCHMARK_CAPTURE(BM_UnpackUint64, Avx2Unaligned, false, &unpack_avx2<uint64_t>,
 #endif
 
 #if defined(ARROW_HAVE_RUNTIME_AVX512)
+BENCHMARK_CAPTURE(BM_UnpackBool, Avx512Unaligned, false, &unpack_avx512<bool>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
+                  "Avx512 not available")
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, Avx512Unaligned, false, &unpack_avx512<uint8_t>,
+                  !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
+                  "Avx512 not available")
+    ->ArgsProduct(kBitWidthsNumValues8);
 BENCHMARK_CAPTURE(BM_UnpackUint16, Avx512Unaligned, false, &unpack_avx512<uint16_t>,
                   !CpuInfo::GetInstance()->IsSupported(CpuInfo::AVX512),
                   "Avx512 not available")
@@ -167,6 +230,10 @@ BENCHMARK_CAPTURE(BM_UnpackUint64, Avx512Unaligned, false, &unpack_avx512<uint64
 #endif
 
 #if defined(ARROW_HAVE_NEON)
+BENCHMARK_CAPTURE(BM_UnpackBool, NeonUnaligned, false, &unpack_neon<bool>)
+    ->ArgsProduct(kBitWidthsNumValuesBool);
+BENCHMARK_CAPTURE(BM_UnpackUint8, NeonUnaligned, false, &unpack_neon<uint8_t>)
+    ->ArgsProduct(kBitWidthsNumValues8);
 BENCHMARK_CAPTURE(BM_UnpackUint16, NeonUnaligned, false, &unpack_neon<uint16_t>)
     ->ArgsProduct(kBitWidthsNumValues16);
 BENCHMARK_CAPTURE(BM_UnpackUint32, NeonUnaligned, false, &unpack_neon<uint32_t>)
@@ -175,6 +242,16 @@ BENCHMARK_CAPTURE(BM_UnpackUint64, NeonUnaligned, false, &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>)