From 59f97b0eb18cb36d48c00bd53d23649c7244d970 Mon Sep 17 00:00:00 2001 From: Jimmy Lu Date: Wed, 8 May 2024 11:32:49 -0700 Subject: [PATCH] [native] Switch to HEAD method for exchange get size call --- .../presto_cpp/main/PrestoExchangeSource.cpp | 16 ++++++++-------- .../presto_cpp/main/TaskResource.cpp | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/presto-native-execution/presto_cpp/main/PrestoExchangeSource.cpp b/presto-native-execution/presto_cpp/main/PrestoExchangeSource.cpp index e510cc89f276..87f085bb860c 100644 --- a/presto-native-execution/presto_cpp/main/PrestoExchangeSource.cpp +++ b/presto-native-execution/presto_cpp/main/PrestoExchangeSource.cpp @@ -179,17 +179,16 @@ void PrestoExchangeSource::doRequest( } auto path = fmt::format("{}/{}", basePath_, sequence_); - VLOG(1) << "Fetching data from " << host_ << ":" << port_ << " " << path; auto self = getSelfPtr(); - auto requestBuilder = - http::RequestBuilder().method(proxygen::HTTPMethod::GET).url(path); - + proxygen::HTTPMethod method; if (maxBytes == 0) { - requestBuilder.header(protocol::PRESTO_GET_DATA_SIZE_HEADER, "true"); - // Coordinator ignores the header and always sends back data. There is only - // one coordinator to fetch data from, so a limit of 1MB is enough. - maxBytes = 1 << 20; + method = proxygen::HTTPMethod::HEAD; + VLOG(1) << "Get data sizes from " << host_ << ":" << port_ << " " << path; + } else { + method = proxygen::HTTPMethod::GET; + VLOG(1) << "Fetching data from " << host_ << ":" << port_ << " " << path; } + auto requestBuilder = http::RequestBuilder().method(method).url(path); velox::common::testutil::TestValue::adjust( "facebook::presto::PrestoExchangeSource::doRequest", this); @@ -289,6 +288,7 @@ void PrestoExchangeSource::processDataResponse( auto remainingBytesString = headers->getHeaders().getSingleOrEmpty( protocol::PRESTO_BUFFER_REMAINING_BYTES_HEADER); if (!remainingBytesString.empty()) { + VLOG(1) << "Remaining bytes: " << remainingBytesString; folly::split(',', remainingBytesString, remainingBytes); if (!remainingBytes.empty() && remainingBytes[0] == 0) { VELOX_CHECK_EQ(remainingBytes.size(), 1); diff --git a/presto-native-execution/presto_cpp/main/TaskResource.cpp b/presto-native-execution/presto_cpp/main/TaskResource.cpp index 232054f84ca6..34b765b87a1a 100644 --- a/presto-native-execution/presto_cpp/main/TaskResource.cpp +++ b/presto-native-execution/presto_cpp/main/TaskResource.cpp @@ -416,7 +416,7 @@ proxygen::RequestHandler* TaskResource::getResults( auto maxWait = getMaxWait(message).value_or( protocol::Duration(protocol::PRESTO_MAX_WAIT_DEFAULT)); protocol::DataSize maxSize; - if (getDataSize || headers.exists(protocol::PRESTO_GET_DATA_SIZE_HEADER)) { + if (getDataSize) { maxSize = protocol::DataSize(0, protocol::DataUnit::BYTE); } else { maxSize = protocol::DataSize(