[native] Switch to HEAD method for exchange get size call

prestodb · May 8, 2024 · 59f97b0 · 59f97b0
1 parent fdb0572
commit 59f97b0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git 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
@@ -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(