diff --git a/src/Processors/Formats/Impl/Parquet/Decoding.cpp b/src/Processors/Formats/Impl/Parquet/Decoding.cpp index ca7b8f5057a35..57beab2629cdf 100644 --- a/src/Processors/Formats/Impl/Parquet/Decoding.cpp +++ b/src/Processors/Formats/Impl/Parquet/Decoding.cpp @@ -87,6 +87,69 @@ struct BitPackedRLEDecoder : public PageDecoder skipOrDecode(num_values, &out[start]); } + bool decodeAndIndex(size_t num_values, Dictionary & dictionary, IColumn & out) override + { + /// For Mode::Column the gather goes through IColumn::index, which needs the indexes as a + /// column anyway; let the caller take the unfused path. + if (dictionary.mode == Dictionary::Mode::Column) + return false; + + if (bit_width == 0) + { + /// All indexes are 0. + if (limit == 0) + throw Exception(ErrorCodes::INCORRECT_DATA, "Dict index or rep/def level out of bounds (rle)"); + dictionary.appendRepeated(0, num_values, out); + return true; + } + + const T value_mask = T((1ul << bit_width) - 1); + while (num_values) + { + if (run_length == 0) + startRun(); + + size_t n = std::min(run_length, num_values); + run_length -= n; + num_values -= n; + + if (run_is_rle) + { + /// `val` was bounds-checked in startRun. + dictionary.appendRepeated(size_t(val), n, out); + } + else + { + /// Unpack a batch of indexes into a stack buffer, bounds-check the whole batch at + /// once (keeping the throw out of the inner loop), then gather the batch. + while (n) + { + constexpr size_t batch = 512; + UInt32 buf[batch]; + size_t m = std::min(n, batch); + size_t max_seen = 0; + for (size_t i = 0; i < m; ++i) + { + size_t x = 0; + memcpy(&x, data + (bit_idx >> 3), 8); + x = (x >> (bit_idx & 7)) & value_mask; + max_seen = std::max(max_seen, x); + buf[i] = static_cast(x); + bit_idx += bit_width; + } + if (max_seen >= limit) + throw Exception(ErrorCodes::INCORRECT_DATA, "Dict index or rep/def level out of bounds (bp)"); + dictionary.appendIndexes(buf, m, out); + n -= m; + } + + if (!run_length) + data += run_bytes; + } + } + return true; + } + void startRun() { UInt64 len = 0; @@ -1356,9 +1419,8 @@ size_t Dictionary::decodedFootprintUpperBound( } template -static void indexImpl(const PaddedPODArray & indexes, std::span data, std::span to) +static void indexImpl(const UInt32 * indexes, size_t size, std::span data, std::span to) { - size_t size = indexes.size(); for (size_t i = 0; i < size; ++i) memcpy(to.data() + i * value_size, data.data() + indexes[i] * value_size, value_size); } @@ -1366,24 +1428,35 @@ static void indexImpl(const PaddedPODArray & indexes, std::span & indexes = indexes_col.getData(); + if (mode == Mode::Column) + { + ColumnPtr temp = col->index(indexes_col, /*limit*/ 0); + out.insertRangeFrom(*temp, 0, indexes.size()); + return; + } + appendIndexes(indexes.data(), indexes.size(), out); +} + +void Dictionary::appendIndexes(const UInt32 * indexes, size_t n, IColumn & out) +{ switch (mode) { case Mode::FixedSize: { - auto to = out.insertRawUninitialized(indexes.size()); - chassert(to.size() == value_size * indexes.size()); + auto to = out.insertRawUninitialized(n); + chassert(to.size() == value_size * n); /// Short variable-length memcpy is very slow compared to a simple mov, so we dispatch /// to specialized loops covering basic int types. switch (value_size) { - case 1: indexImpl<1>(indexes, data, to); break; - case 2: indexImpl<2>(indexes, data, to); break; - case 3: indexImpl<3>(indexes, data, to); break; - case 4: indexImpl<4>(indexes, data, to); break; - case 8: indexImpl<8>(indexes, data, to); break; - case 16: indexImpl<16>(indexes, data, to); break; + case 1: indexImpl<1>(indexes, n, data, to); break; + case 2: indexImpl<2>(indexes, n, data, to); break; + case 3: indexImpl<3>(indexes, n, data, to); break; + case 4: indexImpl<4>(indexes, n, data, to); break; + case 8: indexImpl<8>(indexes, n, data, to); break; + case 16: indexImpl<16>(indexes, n, data, to); break; default: - for (size_t i = 0; i < indexes.size(); ++i) + for (size_t i = 0; i < n; ++i) memcpy(to.data() + i * value_size, data.data() + indexes[i] * value_size, value_size); } break; @@ -1391,9 +1464,10 @@ void Dictionary::index(const ColumnUInt32 & indexes_col, IColumn & out) case Mode::StringPlain: { auto & c = assert_cast(out); - c.reserve(c.size() + indexes.size()); - for (UInt32 idx : indexes) + c.reserve(c.size() + n); + for (size_t i = 0; i < n; ++i) { + UInt32 idx = indexes[i]; size_t start = offsets[ssize_t(idx) - 1] + 4; // offsets[-1] is ok because of padding size_t len = offsets[idx] - start; /// TODO [parquet]: Try optimizing short memcpy by taking advantage of padding (maybe memcpySmall.h helps). Also in PlainStringDecoder. @@ -1402,11 +1476,56 @@ void Dictionary::index(const ColumnUInt32 & indexes_col, IColumn & out) break; } case Mode::Column: + /// Handled in `index` (needs the indexes as a column); the fused decode-and-index path + /// declines Mode::Column before getting here (see PageDecoder::decodeAndIndex). + case Mode::Uninitialized: chassert(false); + } +} + +template +static void fillImpl(const char * src, std::span to, size_t n) +{ + for (size_t i = 0; i < n; ++i) + memcpy(to.data() + i * value_size, src, value_size); +} + +void Dictionary::appendRepeated(size_t idx, size_t n, IColumn & out) +{ + chassert(idx < count); + switch (mode) + { + case Mode::FixedSize: { - ColumnPtr temp = col->index(indexes_col, /*limit*/ 0); - out.insertRangeFrom(*temp, 0, indexes.size()); + auto to = out.insertRawUninitialized(n); + chassert(to.size() == value_size * n); + const char * src = data.data() + idx * value_size; + switch (value_size) + { + case 1: memset(to.data(), *src, n); break; + case 2: fillImpl<2>(src, to, n); break; + case 3: fillImpl<3>(src, to, n); break; + case 4: fillImpl<4>(src, to, n); break; + case 8: fillImpl<8>(src, to, n); break; + case 16: fillImpl<16>(src, to, n); break; + default: + for (size_t i = 0; i < n; ++i) + memcpy(to.data() + i * value_size, src, value_size); + } break; } + case Mode::StringPlain: + { + auto & c = assert_cast(out); + size_t start = offsets[ssize_t(idx) - 1] + 4; // offsets[-1] is ok because of padding + size_t len = offsets[idx] - start; + c.reserve(c.size() + n); + for (size_t i = 0; i < n; ++i) + c.insertData(data.data() + start, len); + break; + } + case Mode::Column: + out.insertManyFrom(*col, idx, n); + break; case Mode::Uninitialized: chassert(false); } } diff --git a/src/Processors/Formats/Impl/Parquet/Decoding.h b/src/Processors/Formats/Impl/Parquet/Decoding.h index e0533fd295073..baf77d24fbaae 100644 --- a/src/Processors/Formats/Impl/Parquet/Decoding.h +++ b/src/Processors/Formats/Impl/Parquet/Decoding.h @@ -55,6 +55,14 @@ struct Dictionary /// decoded `col`), excluding `data` which only points into one of those or into prefetcher memory. size_t allocatedBytes() const; void index(const ColumnUInt32 & indexes_col, IColumn & out); + /// Append the values at the given dictionary indexes to `out`. Same as `index`, from a plain + /// array; `index` delegates here. The indexes must be within bounds. + void appendIndexes(const UInt32 * indexes, size_t n, IColumn & out); + /// Append the value at dictionary index `idx` to `out`, `n` times. Used by the fused + /// decode-and-index path (`PageDecoder::decodeAndIndex`) to turn an RLE run of a repeated + /// index into a bulk fill, instead of expanding the run into explicit indexes and gathering + /// them one by one. + void appendRepeated(size_t idx, size_t n, IColumn & out); void decode(parq::Encoding::type encoding, const PageDecoderInfo & info, size_t num_values, std::span data_, const IDataType & raw_decoded_type); /// Upper bound on `allocatedBytes()` after `decode()` with the given arguments, computed from the @@ -78,6 +86,15 @@ struct PageDecoder virtual void skip(size_t num_values) = 0; virtual void decode(size_t num_values, IColumn & col, const UInt8 * filter, size_t filter_offset) = 0; + /// Fused decode-and-gather for dictionary-encoded data pages: append the *dictionary values* + /// at the decoded indexes directly to `out`, without materializing the indexes as a column. + /// An RLE run of a repeated index becomes a single dictionary lookup and a bulk fill, and + /// bit-packed indexes are gathered from a small stack buffer, saving a full write + read pass + /// over an indexes column (most values in typical files sit in RLE runs). Returns false if + /// this decoder (or this dictionary mode) does not support the fusion; the caller then falls + /// back to `decode` into an indexes column + `Dictionary::index`. + virtual bool decodeAndIndex(size_t /*num_values*/, Dictionary & /*dictionary*/, IColumn & /*out*/) { return false; } + explicit PageDecoder(std::span data_) : data(data_.data()), end(data_.data() + data_.size()) {} virtual ~PageDecoder() = default; diff --git a/src/Processors/Formats/Impl/Parquet/Reader.cpp b/src/Processors/Formats/Impl/Parquet/Reader.cpp index b8c3750c1070d..5279e25c6189b 100644 --- a/src/Processors/Formats/Impl/Parquet/Reader.cpp +++ b/src/Processors/Formats/Impl/Parquet/Reader.cpp @@ -3226,15 +3226,20 @@ void Reader::readRowsInPage(size_t end_row_idx, ColumnSubchunk & subchunk, Colum if (page.is_dictionary_encoded) { - if (!page.indices_column) - page.indices_column = ColumnUInt32::create(); - auto & indices_column_uint32 = assert_cast(*page.indices_column); - auto & data = indices_column_uint32.getData(); - chassert(data.empty()); chassert(!filter); - page.decoder->decode(encoded_values_to_read, *page.indices_column, nullptr, 0); - column.dictionary.index(indices_column_uint32, *subchunk.column); - data.clear(); + /// Fused decode-and-gather; falls back to materializing the indexes as a column when + /// the decoder or the dictionary mode does not support the fusion. + if (!page.decoder->decodeAndIndex(encoded_values_to_read, column.dictionary, *subchunk.column)) + { + if (!page.indices_column) + page.indices_column = ColumnUInt32::create(); + auto & indices_column_uint32 = assert_cast(*page.indices_column); + auto & data = indices_column_uint32.getData(); + chassert(data.empty()); + page.decoder->decode(encoded_values_to_read, *page.indices_column, nullptr, 0); + column.dictionary.index(indices_column_uint32, *subchunk.column); + data.clear(); + } } else {