mapleFU commented on code in PR #55408:
URL: https://github.com/apache/arrow/pull/55408#discussion_r4478586374
##########
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:
I don't know, it would making the very short loops (9-31 elements?) a bit
slower ( like 1.5ns -> 3ns ) on my Mac M4 Pro. If this doesn't matter just
counting by chunk is ok.
--
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]
