HuaHuaY commented on code in PR #55408:
URL: https://github.com/apache/arrow/pull/55408#discussion_r4478437504
##########
cpp/src/parquet/level_comparison.h:
##########
@@ -35,4 +36,78 @@ struct MinMax {
MinMax PARQUET_EXPORT FindMinMax(const int16_t* levels, int64_t num_levels);
+/// Counts the levels that satisfy `pred`, in chunks small enough that a
+/// `uint16_t` accumulator cannot overflow.
+///
+/// This is the fallback for inputs that do not fit in a `uint16_t`
accumulator.
+/// It is kept out of line so that the common path (a single Parquet page
chunk,
+/// always short enough for one accumulator) is a plain loop without any chunk
+/// bookkeeping in front of it.
+template <typename Predicate>
+ARROW_NOINLINE int64_t CountIf16Chunked(const int16_t* levels, int64_t
num_levels,
+ Predicate pred) {
+ // The loop below accumulates one match per lane in a `uint16_t`, so a chunk
+ // must not hold more levels than that can count.
+ constexpr int64_t kMaxCount = std::numeric_limits<uint16_t>::max();
+ int64_t total = 0;
+ while (num_levels > 0) {
+ const int64_t chunk_length = std::min(num_levels, kMaxCount);
+ uint16_t chunk_count = 0;
+ for (int64_t i = 0; i < chunk_length; ++i) {
+ chunk_count += pred(levels[i]) ? 1 : 0;
+ }
+ total += chunk_count;
+ levels += chunk_length;
+ num_levels -= chunk_length;
+ }
+ return total;
+}
+
+/// Count the levels that satisfy the given predicate.
+///
+/// The matches are accumulated in a `uint16_t`, which cannot overflow as long
+/// as at most `std::numeric_limits<uint16_t>::max()` levels are counted at a
+/// time, so longer inputs are delegated to `CountIf16Chunked`. This is both
+/// safe and much faster than `std::count`, as the loop can be vectorized with
+/// 16-bit accumulators (one accumulated match per lane) instead of widening
+/// every match to a 64-bit integer.
+///
+/// \param[in] levels The level values to inspect.
+/// \param[in] num_levels The number of levels to inspect, must be
non-negative.
+/// \param[in] pred The predicate to apply, must be callable as
+/// `bool(int16_t)`.
+/// \return The number of levels for which `pred` returns true.
+template <typename Predicate>
+inline int64_t CountIf16(const int16_t* levels, int64_t num_levels, Predicate
pred) {
+ // Short inputs are counted first, before the wide loop and before the chunk
+ // clamp, because a handful of levels cannot amortize either. Running the
+ // small case first is also how folly::simd::detail::simdForEach orders its
+ // single steps before the unrolled ones, so that short inputs never pay for
+ // the unrolled set up.
+ //
+ // 32 is the widest per-iteration body this reduction gets: a packed int16_t
+ // register holds 8 levels with SSE2/NEON and 16 with AVX2, but clang runs up
+ // to 32 levels per iteration on NEON (four registers), so shorter inputs
+ // never reach a full iteration on the narrow targets. The bound also has to
+ // be a compile-time power of two, which is what turns the loop into an
+ // unrolled ladder of 8/16/24-wide compares instead of scalar code.
+ constexpr int64_t kSmallLevels = 32;
+ constexpr int64_t kMaxCount = std::numeric_limits<uint16_t>::max();
+ if (num_levels < kSmallLevels) {
Review Comment:
Does it make sense to keep this `if` branch, rather than letting LLVM
automatically generate optimized code for short loops?
##########
cpp/src/parquet/size_statistics.cc:
##########
@@ -160,17 +160,11 @@ void UpdateLevelHistogram(std::span<const int16_t>
levels, std::span<int64_t> hi
if (max_level == 1) {
// Specialize the common case for non-repeated non-nested columns.
- // Summing the levels gives us the number of 1s, and the number of 0s
follows.
- // We do repeated sums in the int16_t space, which the compiler is likely
- // to vectorize efficiently.
- constexpr int64_t kChunkSize = 1 << 14; // to avoid int16_t overflows
- int64_t hist1 = 0;
- auto it = levels.begin();
- while (it != levels.end()) {
- const auto chunk_size = std::min<int64_t>(levels.end() - it, kChunkSize);
- hist1 += std::accumulate(it, it + chunk_size, int16_t{0});
- it += chunk_size;
- }
+ // Counting the 1s directly gives us the number of 1s, and the number of 0s
+ // follows. CountIf16 accumulates in the int16_t space in chunks, which the
+ // compiler vectorizes efficiently and which cannot overflow.
+ const int64_t hist1 = internal::CountIf16(levels.data(), num_levels,
Review Comment:
My LLM pointed out a performance regression here, and I think it’s right.
There’s a logical constraint—each `level` can only be 1 or 0, so a summation
could replace `std::count`. However, the compiler can’t infer this, and the new
code ends up performing a check (`level == 1 ? 1 : 0`), which actually makes it
slower.
--
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]
