PERF: Avoid duplicate HEAD requests for file size with S3 endpoint

cpp/src/remote_handle.cpp

@@ -625,6 +625,7 @@ RemoteHandle RemoteHandle::open(std::string url,
   };
 
   std::unique_ptr<RemoteEndpoint> endpoint;
+  std::optional<std::size_t> probed_nbytes;
 
   if (remote_endpoint_type == RemoteEndpointType::AUTO) {
     // Try each allowed type in the order of allowlist
@@ -634,16 +635,17 @@ RemoteHandle RemoteHandle::open(std::string url,
         if (endpoint == nullptr) { continue; }
         if (type == RemoteEndpointType::S3) {
           // Check connectivity for the credential-based S3 endpoint, and throw an exception if
-          // failed
-          endpoint->get_file_size();
+          // failed. Reuse this size when constructing the handle to avoid a second HEAD request.
+          probed_nbytes = endpoint->get_file_size();
         }
       } catch (...) {
         // If the credential-based S3 endpoint cannot be used to access the URL, try using S3 public
         // endpoint instead if it is in the allowlist
         if (type == RemoteEndpointType::S3 &&
             std::find(allow_list->begin(), allow_list->end(), RemoteEndpointType::S3_PUBLIC) !=
               allow_list->end()) {
-          endpoint = std::make_unique<S3PublicEndpoint>(url);
+          endpoint      = std::make_unique<S3PublicEndpoint>(url);
+          probed_nbytes = std::nullopt;
         } else {
           throw;
         }
@@ -669,8 +671,11 @@ RemoteHandle RemoteHandle::open(std::string url,
                   std::runtime_error);
   }
 
-  return nbytes.has_value() ? RemoteHandle(std::move(endpoint), nbytes.value())
-                            : RemoteHandle(std::move(endpoint));
+  if (nbytes.has_value()) { return RemoteHandle(std::move(endpoint), nbytes.value()); }
+  if (probed_nbytes.has_value()) {
+    return RemoteHandle(std::move(endpoint), probed_nbytes.value());
+  }
+  return RemoteHandle(std::move(endpoint));
 }
 
 RemoteHandle::RemoteHandle(std::unique_ptr<RemoteEndpoint> endpoint, std::size_t nbytes)

cpp/tests/test_remote_handle.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -67,12 +67,12 @@ class RemoteHandleTest : public testing::Test {
         // Test unified interface
         {
           // Here we pass the 1-byte argument to RemoteHandle::open. For all endpoints except
-          // kvikio::RemoteEndpointType::S3, this prevents the endpoint constructor from querying
-          // the file size and sending requests to the server, thus allowing us to use dummy URLs
-          // for testing purpose.
-          // For kvikio::RemoteEndpointType::S3, RemoteHandle::open sends HEAD request as a
-          // connectivity check and will fail on the syntactically valid dummy URL. The
-          // kvikio::RemoteEndpointType::S3_PUBLIC will then be used as the endpoint.
+          // kvikio::RemoteEndpointType::S3 in AUTO mode, this prevents querying the file size and
+          // sending requests to the server, thus allowing us to use dummy URLs for testing.
+          // For kvikio::RemoteEndpointType::S3 with AUTO, RemoteHandle::open sends a HEAD request
+          // as a connectivity check (and reuses that size when nbytes is not provided). It will
+          // fail on the syntactically valid dummy URL, and kvikio::RemoteEndpointType::S3_PUBLIC
+          // will then be used as the endpoint.
           auto remote_handle =
             kvikio::RemoteHandle::open(url, kvikio::RemoteEndpointType::AUTO, std::nullopt, 1);
           EXPECT_EQ(remote_handle.remote_endpoint_type(), expected_endpoint_type);