kecookier commented on code in PR #12984:
URL: https://github.com/apache/gluten/pull/12984#discussion_r4047004537
##########
cpp/velox/shuffle/VeloxShuffleReader.cc:
##########
@@ -948,13 +1011,173 @@ void
VeloxRssSortShuffleReaderDeserializer::loadNextStream() {
constexpr uint64_t kMaxReadBufferSize = (1 << 20) -
AlignedBuffer::kPaddedSize;
auto buffer = AlignedBuffer::allocate<char>(kMaxReadBufferSize,
memoryManager_->getLeafMemoryPool().get());
- in_ = std::make_unique<VeloxInputStream>(std::move(arrowIn_),
std::move(buffer));
+ in_ = std::make_unique<RssSortShuffleReaderInputStream>(std::move(arrowIn_),
std::move(buffer));
}
-size_t
VeloxRssSortShuffleReaderDeserializer::VeloxInputStream::remainingSize() const {
+size_t
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::remainingSize()
const {
return std::numeric_limits<unsigned long>::max();
}
+uint8_t
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::readByte()
{
+ if (current_->position < current_->size) {
+ return current_->buffer[current_->position++];
+ }
+ next();
+ return readByte();
+}
+
+void
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::readBytes(uint8_t*
bytes, int32_t size) {
+ VELOX_CHECK_GE(size, 0, "Attempting to read negative number of bytes");
+ int32_t offset = 0;
+ for (;;) {
+ int32_t available = current_->size - current_->position;
+ int32_t numUsed = std::min(available, size);
+ simd::memcpy(bytes + offset, current_->buffer + current_->position,
numUsed);
+ offset += numUsed;
+ size -= numUsed;
+ current_->position += numUsed;
+ if (!size) {
+ return;
+ }
+ next();
+ }
+}
+
+void
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::skip(int32_t
size) {
+ VELOX_CHECK_GE(size, 0, "Attempting to skip negative number of bytes");
+ for (;;) {
+ int32_t available = current_->size - current_->position;
+ int32_t numUsed = std::min(available, size);
+ size -= numUsed;
+ current_->position += numUsed;
+ if (!size) {
+ return;
+ }
+ next();
+ }
+}
+
+std::string
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::toString()
const {
+ std::stringstream oss;
+ oss << ranges_.size() << " ranges (position/size) [";
+ for (const auto& range : ranges_) {
+ oss << "(" << range.position << "/" << range.size << (&range == current_ ?
" current" : "") << ")";
+ if (&range != &ranges_.back()) {
+ oss << ",";
+ }
+ }
+ oss << "]";
+ return oss.str();
+}
+
+std::string_view
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::nextView(int64_t
size) {
+ VELOX_CHECK_GE(size, 0, "Attempting to view negative number of bytes");
+ if (ranges_.empty()) {
+ return std::string_view(nullptr, 0);
+ }
+ if (ranges_[0].position == ranges_[0].size) {
+ // Current window is exhausted. For single-window refill streams, next()
+ // overwrites the window with fresh data, so attempt refill before
+ // reporting end-of-stream.
+ next(false);
+ if (ranges_.empty() || ranges_[0].position == ranges_[0].size) {
+ return std::string_view(nullptr, 0);
+ }
+ }
+ VELOX_DCHECK(ranges_[0].size > 0);
+ const int32_t position = ranges_[0].position;
+ const int64_t viewSize = std::min<int64_t>(ranges_[0].size -
ranges_[0].position, size);
+ ranges_[0].position += static_cast<int32_t>(viewSize);
+ return std::string_view(reinterpret_cast<char*>(ranges_[0].buffer) +
position, viewSize);
+}
+
+std::streampos
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::tellp()
const {
+ if (ranges_.empty()) {
+ return 0;
+ }
+ return static_cast<std::streampos>(static_cast<int64_t>(totalBytesRead_) -
(ranges_[0].size - ranges_[0].position));
+}
+
+void
VeloxRssSortShuffleReaderDeserializer::RssSortShuffleReaderInputStream::seekp(std::streampos
position) {
+ if (ranges_.empty() && position == 0) {
+ return;
+ }
+ VELOX_CHECK(!ranges_.empty(), "Cannot seek an empty
RssSortShuffleReaderInputStream");
+ const int64_t windowStart = static_cast<int64_t>(totalBytesRead_) -
ranges_[0].size;
+ const int64_t windowEnd = static_cast<int64_t>(totalBytesRead_);
+ const int64_t target = static_cast<int64_t>(position);
+ VELOX_CHECK(
+ target >= windowStart && target <= windowEnd,
+ "RssSortShuffleReaderInputStream::seekp({}) is outside the resident
window [{}, {}): bytes before the "
+ "window were already consumed from the underlying stream
(totalBytesRead={})",
+ target,
+ windowStart,
+ windowEnd,
+ totalBytesRead_);
+ ranges_[0].position = static_cast<int32_t>(target - windowStart);
+}
+
+RowVectorPtr VeloxRssSortShuffleReaderDeserializer::readPage() {
+ using facebook::velox::serializer::presto::detail::kCompressedBitMask;
+ using facebook::velox::serializer::presto::detail::kHeaderSize;
+ using facebook::velox::serializer::presto::detail::PrestoHeader;
+ constexpr int32_t kPrestoHeaderSize = kHeaderSize;
+
+ // Fast path: peek the header without consuming; if the whole page fits in
+ // the current window, deserialize in-situ from in_ directly.
+ if (in_->remainingInWindow() >= kPrestoHeaderSize) {
+ std::string_view window(reinterpret_cast<const char*>(in_->data()),
in_->remainingInWindow());
+ auto peekedHeader = PrestoHeader::read(&window);
+ if (peekedHeader.has_value()) {
+ const int32_t payloadSize = (peekedHeader->pageCodecMarker &
kCompressedBitMask) != 0
+ ? peekedHeader->compressedSize
+ : peekedHeader->uncompressedSize;
+ const int64_t totalSize = kPrestoHeaderSize +
static_cast<int64_t>(payloadSize);
+ if (totalSize <= in_->remainingInWindow()) {
+ RowVectorPtr rowVector;
+ VectorStreamGroup::read(
+ in_.get(), memoryManager_->getLeafMemoryPool().get(), rowType_,
serde_, &rowVector, &serdeOptions_);
+ return rowVector;
+ }
+ }
+ }
+
+ // Slow path: the page spans multiple read windows. Reassemble it into a
+ // contiguous BufferInputStream so the serde's backward seek never touches
+ // window data overwritten by a refill.
+ std::array<uint8_t, kHeaderSize> headerStorage;
+ in_->readBytes(headerStorage.data(), kHeaderSize);
+ std::string_view headerBytes(reinterpret_cast<const
char*>(headerStorage.data()), kHeaderSize);
+ auto headerOpt = PrestoHeader::read(&headerBytes);
+ VELOX_CHECK(headerOpt.has_value(), "Invalid Presto page header");
+ const auto& header = *headerOpt;
+
+ const int32_t payloadSize =
+ (header.pageCodecMarker & kCompressedBitMask) != 0 ?
header.compressedSize : header.uncompressedSize;
Review Comment:
Not needed here.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
