PERF/API: Specify datasource size to avoid HEAD requests for S3 endpoints

cpp/include/cudf/io/datasource.hpp

@@ -13,7 +13,7 @@
 #include <rmm/cuda_stream_view.hpp>
 
 #include <future>
-#include <memory>
+#include <optional>
 
 namespace CUDF_EXPORT cudf {
 //! IO interfaces
@@ -98,11 +98,14 @@ class datasource {
    * @param[in] offset Starting byte offset from which data will be read (the default is zero)
    * @param[in] max_size_estimate Upper estimate of the data range that will be read (the default is
    * zero, which means the whole file after `offset`)
+   * @param[in] known_size Optional known file size in bytes. When set for remote URLs, KvikIO skips
+   * the size query HEAD request at open time.
    * @return Constructed datasource object
    */
   static std::unique_ptr<datasource> create(std::string const& filepath,
-                                            size_t offset            = 0,
-                                            size_t max_size_estimate = 0);
+                                            size_t offset                         = 0,
+                                            size_t max_size_estimate              = 0,
+                                            std::optional<std::size_t> known_size = std::nullopt);
 
   /**
    * @brief Creates a source from a host memory buffer.

cpp/include/cudf/io/types.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -310,6 +310,17 @@ constexpr inline auto is_byte_like_type()
          std::is_same_v<non_cv_T, std::byte>;
 }
 
+/**
+ * @brief A file path with an optional known size in bytes.
+ *
+ * When `size` is set for a remote URL, libcudf passes it to KvikIO at open time so the remote
+ * server is not queried for file size.
+ */
+struct filepath_source {
+  std::string path;                   ///< Path or URL of the input file
+  std::optional<std::size_t> size{};  ///< Known file size; omit to query via KvikIO
+};
+
 /**
  * @brief Source information for read interfaces
  */
@@ -325,8 +336,13 @@ struct source_info {
    * @param file_paths Input files paths
    */
   explicit source_info(std::vector<std::string> file_paths)
-    : _type(io_type::FILEPATH), _num_sources(file_paths.size()), _filepaths(std::move(file_paths))
+    : _type(io_type::FILEPATH), _num_sources(file_paths.size())
   {
+    _filepath_sources.reserve(file_paths.size());
+    for (auto& path : file_paths) {
+      _filepath_sources.push_back({std::move(path), std::nullopt});
+    }
+    rebuild_filepaths();
   }
 
   /**
@@ -335,10 +351,21 @@ struct source_info {
    * @param file_path Single input file
    */
   explicit source_info(std::string file_path)
-    : _type(io_type::FILEPATH), _num_sources(1), _filepaths({std::move(file_path)})
+    : source_info(std::vector<std::string>{std::move(file_path)})
   {
   }
 
+  /**
+   * @brief Construct a new source info object from filepath sources with optional known sizes
+   *
+   * @param sources Input filepath sources
+   */
+  explicit source_info(std::vector<filepath_source> sources)
+    : _type(io_type::FILEPATH), _num_sources(sources.size()), _filepath_sources(std::move(sources))
+  {
+    rebuild_filepaths();
+  }
+
   /**
    * @brief Construct a new source info object for multiple buffers in host memory
    *
@@ -424,6 +451,12 @@ struct source_info {
    * @return The type of the input
    */
   [[nodiscard]] auto type() const { return _type; }
+  /**
+   * @brief Get the filepath sources of the input
+   *
+   * @return The filepath sources of the input
+   */
+  [[nodiscard]] auto const& filepath_sources() const { return _filepath_sources; }
   /**
    * @brief Get the filepaths of the input
    *
@@ -457,8 +490,18 @@ struct source_info {
   [[nodiscard]] auto num_sources() const { return _num_sources; }
 
  private:
+  void rebuild_filepaths()
+  {
+    _filepaths.clear();
+    _filepaths.reserve(_filepath_sources.size());
+    for (auto const& source : _filepath_sources) {
+      _filepaths.push_back(source.path);
+    }
+  }
+
   io_type _type       = io_type::VOID;
   size_t _num_sources = 0;
+  std::vector<filepath_source> _filepath_sources;
   std::vector<std::string> _filepaths;
   std::vector<cudf::host_span<std::byte const>> _host_buffers;
   std::vector<cudf::device_span<std::byte const>> _device_buffers;

cpp/src/io/functions.cpp

@@ -162,24 +162,26 @@ std::vector<std::unique_ptr<cudf::io::datasource>> make_datasources(source_info
   switch (info.type()) {
     case io_type::FILEPATH: {
       std::vector<std::unique_ptr<cudf::io::datasource>> sources;
-      sources.reserve(info.filepaths().size());
+      sources.reserve(info.filepath_sources().size());
       // Creating sources in a single thread is faster for a small number of sources
       auto const pool_use_threshold =
         cudf::detail::getenv_or("LIBCUDF_DATASOURCE_PARALLEL_CREATION_THRESHOLD", 8ul);
-      if (info.filepaths().size() >= pool_use_threshold) {
+      if (info.filepath_sources().size() >= pool_use_threshold) {
         std::vector<std::future<std::unique_ptr<cudf::io::datasource>>> source_tasks;
-        source_tasks.reserve(info.filepaths().size());
-        for (auto const& path : info.filepaths()) {
-          source_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task(
-            [=] { return cudf::io::datasource::create(path, offset, max_size_estimate); }));
+        source_tasks.reserve(info.filepath_sources().size());
+        for (auto const& fs : info.filepath_sources()) {
+          source_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task([=] {
+            return cudf::io::datasource::create(fs.path, offset, max_size_estimate, fs.size);
+          }));
         }
         std::transform(
           source_tasks.begin(), source_tasks.end(), std::back_inserter(sources), [](auto& task) {
             return task.get();
           });
       } else {
-        for (auto const& filepath : info.filepaths()) {
-          sources.emplace_back(cudf::io::datasource::create(filepath, offset, max_size_estimate));
+        for (auto const& fs : info.filepath_sources()) {
+          sources.emplace_back(
+            cudf::io::datasource::create(fs.path, offset, max_size_estimate, fs.size));
         }
       }
       return sources;

cpp/src/io/utilities/datasource.cpp

@@ -27,6 +27,7 @@
 #include <vector>
 
 #ifdef CUDF_KVIKIO_REMOTE_IO
+#include <kvikio/hdfs.hpp>
 #include <kvikio/remote_handle.hpp>
 #endif
 
@@ -353,14 +354,43 @@ class user_datasource_wrapper : public datasource {
   datasource* const source;  ///< A non-owning pointer to the user-implemented datasource
 };
 
+#ifdef CUDF_KVIKIO_REMOTE_IO
+/**
+ * @brief Infer the KvikIO remote endpoint type from a URL (no network I/O).
+ *
+ * Mirrors the order used by `kvikio::RemoteHandle::open()` in AUTO mode.
+ */
+kvikio::RemoteEndpointType infer_remote_endpoint_type(std::string const& url)
+{
+  if (kvikio::S3Endpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::S3; }
+  if (kvikio::S3PublicEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::S3_PUBLIC; }
+  if (kvikio::S3EndpointWithPresignedUrl::is_url_valid(url)) {
+    return kvikio::RemoteEndpointType::S3_PRESIGNED_URL;
+  }
+  if (kvikio::WebHdfsEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::WEBHDFS; }
+  if (kvikio::HttpEndpoint::is_url_valid(url)) { return kvikio::RemoteEndpointType::HTTP; }
+  return kvikio::RemoteEndpointType::HTTP;
+}
+
+kvikio::RemoteHandle open_remote_handle(char const* filepath, std::optional<std::size_t> known_size)
+{
+  if (known_size.has_value()) {
+    auto const endpoint_type = infer_remote_endpoint_type(filepath);
+    return kvikio::RemoteHandle::open(filepath, endpoint_type, std::nullopt, *known_size);
+  }
+  return kvikio::RemoteHandle::open(filepath);
+}
+#endif
+
 #ifdef CUDF_KVIKIO_REMOTE_IO
 /**
  * @brief Remote file source backed by KvikIO, which handles S3 filepaths seamlessly.
  */
 class remote_file_source : public kvikio_source<kvikio::RemoteHandle> {
  public:
-  explicit remote_file_source(char const* filepath)
-    : kvikio_source{kvikio::RemoteHandle::open(filepath)}
+  explicit remote_file_source(char const* filepath,
+                              std::optional<std::size_t> known_size = std::nullopt)
+    : kvikio_source{open_remote_handle(filepath, known_size)}
   {
   }
 
@@ -397,7 +427,8 @@ class remote_file_source : public file_source {
 
 std::unique_ptr<datasource> datasource::create(std::string const& filepath,
                                                size_t offset,
-                                               size_t max_size_estimate)
+                                               size_t max_size_estimate,
+                                               std::optional<std::size_t> known_size)
 {
   auto const use_memory_mapping = [] {
     auto const policy = cudf::detail::getenv_or("LIBCUDF_MMAP_ENABLED", std::string{"OFF"});
@@ -410,7 +441,7 @@ std::unique_ptr<datasource> datasource::create(std::string const& filepath,
 
   if (remote_file_source::could_be_remote_url(filepath)) {
     try {
-      return std::make_unique<remote_file_source>(filepath.c_str());
+      return std::make_unique<remote_file_source>(filepath.c_str(), known_size);
     } catch (std::exception const& ex) {
       std::string redacted_msg;
       try {
@@ -450,7 +481,7 @@ std::unique_ptr<datasource> datasource::create(std::string const& filepath,
       // Create a remote file resource only when the pattern is found and replaced; otherwise, still
       // create a local file resource
       if (filepath != remote_file_path) {
-        return std::make_unique<remote_file_source>(remote_file_path.c_str());
+        return std::make_unique<remote_file_source>(remote_file_path.c_str(), known_size);
       }
     }
 

cpp/tests/CMakeLists.txt

@@ -320,6 +320,7 @@ ConfigureTest(
 # * io tests --------------------------------------------------------------------------------------
 ConfigureTest(COMPRESSION_TEST io/comp/comp_test.cpp)
 ConfigureTest(ROW_SELECTION_TEST io/row_selection_test.cpp)
+ConfigureTest(FILEPATH_SOURCE_TEST io/filepath_source_test.cpp)
 
 ConfigureTest(
   CSV_TEST io/csv_test.cpp

cpp/tests/io/filepath_source_test.cpp

@@ -0,0 +1,73 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_utilities.hpp>
+#include <cudf_test/column_wrapper.hpp>
+
+#include <cudf/io/datasource.hpp>
+#include <cudf/io/parquet.hpp>
+#include <cudf/io/types.hpp>
+
+#include <filesystem>
+
+struct FilepathSourceTest : public cudf::test::BaseFixture {};
+
+TEST_F(FilepathSourceTest, StringConstructorsPopulateFilepathSources)
+{
+  auto const single = cudf::io::source_info{"test.parquet"};
+  ASSERT_EQ(single.filepath_sources().size(), 1);
+  EXPECT_EQ(single.filepath_sources().front().path, "test.parquet");
+  EXPECT_FALSE(single.filepath_sources().front().size.has_value());
+  EXPECT_EQ(single.filepaths(), std::vector<std::string>{"test.parquet"});
+
+  auto const multi = cudf::io::source_info{std::vector<std::string>{"a.parquet", "b.parquet"}};
+  ASSERT_EQ(multi.filepath_sources().size(), 2);
+  EXPECT_EQ(multi.filepaths().size(), 2);
+  EXPECT_FALSE(multi.filepath_sources()[1].size.has_value());
+}
+
+TEST_F(FilepathSourceTest, FilepathSourceConstructorPreservesSize)
+{
+  std::vector<cudf::io::filepath_source> sources{
+    {"s3://bucket/object.parquet", 12345},
+    {"https://example.com/data.parquet", std::nullopt},
+  };
+
+  auto const info = cudf::io::source_info{std::move(sources)};
+  ASSERT_EQ(info.filepath_sources().size(), 2);
+  EXPECT_EQ(info.filepath_sources()[0].path, "s3://bucket/object.parquet");
+  ASSERT_TRUE(info.filepath_sources()[0].size.has_value());
+  EXPECT_EQ(info.filepath_sources()[0].size.value(), 12345);
+  EXPECT_FALSE(info.filepath_sources()[1].size.has_value());
+  EXPECT_EQ(info.filepaths()[0], "s3://bucket/object.parquet");
+  EXPECT_EQ(info.filepaths()[1], "https://example.com/data.parquet");
+}
+
+TEST_F(FilepathSourceTest, KnownSizePlumbsThroughMakeDatasources)
+{
+  auto const filepath =
+    (std::filesystem::temp_directory_path() / "filepath_source_test.parquet").string();
+
+  auto col = cudf::test::fixed_width_column_wrapper<int32_t>{1, 2, 3};
+  cudf::table_view const table{{col}};
+
+  cudf::io::parquet_writer_options write_opts =
+    cudf::io::parquet_writer_options::builder(cudf::io::sink_info{filepath}, table);
+  cudf::io::write_parquet(write_opts);
+
+  auto const file_size = std::filesystem::file_size(filepath);
+  std::vector<cudf::io::filepath_source> sources{{filepath, file_size}};
+  auto const source_info = cudf::io::source_info{std::move(sources)};
+
+  auto datasources = cudf::io::make_datasources(source_info);
+  ASSERT_EQ(datasources.size(), 1);
+  EXPECT_EQ(datasources.front()->size(), file_size);
+
+  auto const read_opts =
+    cudf::io::parquet_reader_options::builder(source_info).column_names({"0"}).build();
+  auto const result = cudf::io::read_parquet(read_opts);
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(col, result.tbl->get_column(0).view());
+}

docs/cudf/source/cudf/io/io.md

@@ -117,6 +117,29 @@ Note that:
 For more information about error handling, compatibility mode, and
 tuning parameters in KvikIO see: <https://github.com/rapidsai/kvikio>
 
+### Remote file sizes and HEAD requests
+
+When reading remote files (for example `s3://...` URLs) via `pylibcudf.io.SourceInfo`, KvikIO
+may send HEAD requests at open time to probe connectivity and query file size. To skip those
+requests when the file size is already known (for example from object-store metadata), pass a
+`pylibcudf.io.FilepathSource` with the `size` argument set:
+
+```python
+import pylibcudf as plc
+
+content_length = ...  # from external metadata
+sources = plc.io.SourceInfo([
+    plc.io.FilepathSource("s3://bucket/object.parquet", size=content_length),
+])
+table = plc.io.parquet.read_parquet(
+    plc.io.parquet.ParquetReaderOptions.builder(sources).build()
+)
+```
+
+Providing an incorrect size avoids the extra HEAD requests but will break footer reads and other
+operations that depend on the true file length. Plain string paths in `SourceInfo` preserve the
+previous behavior (size queried via KvikIO).
+
 Operations that support the use of GPUDirect Storage:
 
 - {py:func}`cudf.read_avro`

python/pylibcudf/pylibcudf/io/__init__.py

@@ -15,10 +15,11 @@
     types,
 )
 from .parquet_metadata import FileMetaData
-from .types import SinkInfo, SourceInfo, TableWithMetadata
+from .types import FilepathSource, SinkInfo, SourceInfo, TableWithMetadata
 
 __all__ = [
     "FileMetaData",
+    "FilepathSource",
     "SinkInfo",
     "SourceInfo",
     "TableWithMetadata",

python/pylibcudf/pylibcudf/io/types.pxd

@@ -14,6 +14,7 @@ from pylibcudf.libcudf.io.types cimport (
     column_name_info,
     compression_type,
     dictionary_policy,
+    filepath_source,
     io_type,
     partition_info,
     quote_style,
@@ -88,6 +89,10 @@ cdef class TableWithMetadata:
         table_with_metadata& tbl, object stream, DeviceMemoryResource mr
     )
 
+cdef class FilepathSource:
+    cdef public object path
+    cdef public object size
+
 cdef class SourceInfo:
     cdef source_info c_obj
     # Keep the bytes converted from stringio alive