Skip to content

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

Merged
rapids-bot[bot] merged 2 commits into
rapidsai:mainfrom
TomAugspurger:tom/kvikio-duplicate-head-requests
Jun 2, 2026
Merged

PERF: Avoid duplicate HEAD requests for file size with S3 endpoint#974
rapids-bot[bot] merged 2 commits into
rapidsai:mainfrom
TomAugspurger:tom/kvikio-duplicate-head-requests

Conversation

@TomAugspurger

@TomAugspurger TomAugspurger commented Jun 1, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Our S3 implementaiton made two HEAD requests on S3 urls in "AUTO" mode (the only mode that's exposed to libcudf users).

  1. A first HEAD request to test S3 credentials
  2. A second HEAD request to get the file size

We can avoid the second HEAD request by reusing the response to the first HEAD request when we go to actually open the file.

Closes #973

Here's an screenshot of an nsys profile showing the single get_file_size() / HEAD request on this script:

import pylibcudf as plc


sources = plc.io.types.SourceInfo(["s3://rapids-tpch/scale-10/nation/part.0.parquet"])
options = plc.io.parquet.ParquetReaderOptions.builder(sources).build()
table = plc.io.parquet.read_parquet(options)
image

@TomAugspurger
PERF: Avoid duplicate HEAD requests for file size with S3 endpoint
b386e6b 
Our S3 implementaiton made two HEAD requests on S3 urls in "AUTO" mode
(the only mode that's exposed to libcudf users).

1. A first HEAD request to test S3 credentials
2. A second HEAD request to get the file size

We can avoid the second HEAD request by reusing the response to the
first HEAD request when we go to actually open the file.
@TomAugspurger
TomAugspurger requested a review from a team as a code owner June 1, 2026 13:50
@TomAugspurger
TomAugspurger marked this pull request as draft June 1, 2026 13:50
@copy-pr-bot

copy-pr-bot Bot commented Jun 1, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@TomAugspurger TomAugspurger mentioned this pull request Jun 1, 2026
Closed
@TomAugspurger TomAugspurger added improvement Improves an existing functionality non-breaking Introduces a non-breaking change labels Jun 1, 2026
@TomAugspurger
TomAugspurger marked this pull request as ready for review June 1, 2026 14:12
@TomAugspurger

TomAugspurger commented Jun 1, 2026

Copy link
Copy Markdown
Contributor Author

I'm looking into the final HTTP head request. This comes from the AUTO endpoint type, which is the only thing used by libcudf, to check whether or not we should fall back to the public endpoint thing. I suspect this will be best solved in libcudf, but I'm not quite sure on that.

This was referenced Jun 1, 2026
Merged
Merged
madsbk
madsbk approved these changes Jun 2, 2026
View reviewed changes

@madsbk madsbk left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good

@TomAugspurger

TomAugspurger commented Jun 2, 2026

Copy link
Copy Markdown
Contributor Author

/merge

@rapids-bot
rapids-bot Bot merged commit 1012701 into rapidsai:main Jun 2, 2026
180 of 184 checks passed
@TomAugspurger
TomAugspurger deleted the tom/kvikio-duplicate-head-requests branch June 2, 2026 14:09
rapids-bot Bot pushed a commit to NVIDIA/cudf that referenced this pull request Jun 28, 2026
@TomAugspurger
PERF/API: Specify datasource size to avoid HEAD requests for S3 endpo…
4d2410a 
…ints (#22739)

Currently, reading, say, parquet metadata or data from a remote source with kvikio requires several HTTP requests. See #22667 (comment) and rapidsai/kvikio#974 for more, but a couple of these HTTP requests are `HEAD` requests used to get the size of the file.

This updates libcudf to provide a way to bypass those `HEAD` requests when the file size is known in advance (#22734 is an example case where we already know the file sizes thanks to a `list` operation). To do this, we expand `SourceInfo` to accept a (list of) simple `plc.io.types.FilepathSource(url, size=content_length)` objects.

For example, `read_parquet_metadata` both without and with file sizes:

<img width="942" height="438" alt="image" src="https://github.com/user-attachments/assets/1f14237c-66b4-4c29-a03d-c5d57c735ba7" />

and `read_parquet`:

<img width="1592" height="441" alt="image" src="https://github.com/user-attachments/assets/06dfcde0-8c3c-4cb7-a37e-d0c9d66efc46" />

Note the lack of `kvikio::RemoteHandle::open` calls in the kvikio domain under the `with-size` annotations. These came from the following script:

<details>

```python
import boto3
import nvtx
import pylibcudf as plc

bucket = "rapids-tpch"
key = "scale-10/nation/part.0.parquet"
url = f"s3://{bucket}/{key}"

with nvtx.annotate("boto3-head_object"):
    response = boto3.client("s3").head_object(Bucket=bucket, Key=key)
    content_length = response["ContentLength"]

sources = plc.io.types.SourceInfo([url])
sources_with_size = plc.io.types.SourceInfo(
    [plc.io.types.FilepathSource(url, size=content_length)]
)


with nvtx.annotate("warmup"):
    plc.io.parquet_metadata.read_parquet_metadata(sources)


with nvtx.annotate("read_parquet_metadata"):
    with nvtx.annotate("no-size"):
        _ = plc.io.parquet_metadata.read_parquet_metadata(sources)

    with nvtx.annotate("with-size"):
        _ = plc.io.parquet_metadata.read_parquet_metadata(sources_with_size)

with nvtx.annotate("read_parquet_footers"):
    with nvtx.annotate("no-size"):
        footers = plc.io.parquet_metadata.read_parquet_footers(sources)

    with nvtx.annotate("with-size"):
        footers_with_size = plc.io.parquet_metadata.read_parquet_footers(sources_with_size)

options = plc.io.parquet.ParquetReaderOptions.builder(sources).build()
options_with_size = plc.io.parquet.ParquetReaderOptions.builder(
    sources_with_size
).build()

with nvtx.annotate("read_parquet"):
    with nvtx.annotate("no-size"):
        table = plc.io.parquet.read_parquet(options)

    with nvtx.annotate("with-size"):
        table = plc.io.parquet.read_parquet(options_with_size)

with nvtx.annotate("read_parquet-with-footers"):
    with nvtx.annotate("no-size"):
        table = plc.io.parquet.read_parquet(options, parquet_metadatas=footers)

    with nvtx.annotate("with-size"):
        table = plc.io.parquet.read_parquet(
            options_with_size, parquet_metadatas=footers_with_size
        )

print("done")

</details>

```

This is a necessary precursor to #22734.

Closes #22740

Authors:
  - Tom Augspurger (https://github.com/TomAugspurger)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Tianyu Liu (https://github.com/kingcrimsontianyu)

URL: #22739
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@madsbk madsbk madsbk approved these changes

Assignees

No one assigned

Labels

improvement Improves an existing functionality non-breaking Introduces a non-breaking change

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Duplicateget_file_size / HEAD request for S3 endpoints

2 participants

@TomAugspurger @madsbk