feat(rust,python): Enable object store in scan_parquet python

[winding-lines]

This PR enables the object_store downloads in python when using scan_parquet. This is opt-in, and "object_store":True needs to be passed in the storage_options parameter.

For example

import polars
from time import time


aws_url = "s3://your-bucket/polars/datasets/foods1.parquet"
start = time()
df = polars.scan_parquet(
    aws_url,
    storage_options={
        "object_store": True,
        "region": "us-west-2",
        "aws_access_key_id": "xxxx",
        "aws_secret_access_key": "yyyy",
    },
)
print(df.head())
print(f"Time: {time() - start}")

[ritchie46]

Thanks a lot for the PR @winding-lines. I hope to get to this one tomorrow. 💪

[winding-lines]

Sorry @ritchie46, can you wait until the weekend/next week?

The code is working but the speedup is not significant - 6 seconds faster on this 42 second test that I am using.

It may well be that the speedup will not be significant when done and that the only benefit is that we can make the planner more efficient, but I want to push a bit more.

[winding-lines]

@ritchie46 here is an update on this PR:

1. this PR is working as described in the python snippet above. I think it is worthwhile for you to review it and based on your feedback we can then merge it.
2. the current Python implementation of scan_parquet("s3://../*.parquet") uses fsspec.open(). In this case fsspec just opens the first file and the subsequent planning phase in polars is also fast. Fast is good but for most users opening just one file is unexpected behavior.
3. the proposed Rust implementation in this PR actually lists all the files and then fetches the file_info in the planning phase.
   a. In my first implementation the async code was more self-contained but this leads to serializing the downloads for the metadata.
   b. In my current implementation of this PR I bring together the listing and downloading phases which cuts about 6 seconds from 42 seconds on a test with 80 parquet files.
4. This PR enables use cases with low tens of parquet files. It is not usable with hundreds of files or more which I think we should also enable in subsequent PRs.

This Dask thread dask/community#234 (comment) has some good pointers on what we could do next. The general idea is to load the metadata from a centralized system, and for a general library as polars it may make sense to support multiple mechanisms. The first one could be the Hive _common_metadata describe in the link.

There is a secondary long term design question: assuming that 1 year from now polars supports Hive, Iceberg, Delta Table, Hudi, where should all of those implementation go? My personal preference is to have them supported natively but I know that adding dependencies to Polars is also not desirable.

[talawahtech]

Great work @winding-lines! Some unsolicited feedback on the function signature, what do you think of making it a little more generic e.g.

df = polars.scan_parquet(
    aws_url,
    storage_options={
        "storage_handler": "object_store",
        "region": "us-west-2",
        "aws_access_key_id": "xxxx",
        "aws_secret_access_key": "yyyy",
    },
)

or

df = polars.scan_parquet(
    aws_url,
    storage_handler="object_store",
    storage_options={
        "region": "us-west-2",
        "aws_access_key_id": "xxxx",
        "aws_secret_access_key": "yyyy",
    },
)

If the storage_handler (or maybe storage_backend) param is missing, it defaults to fsspec.

[kylebarron]

The first one could be the Hive _common_metadata describe in the link.

It's not a part of the parquet spec, but systems that write directories of Parquet files often produce both _common_metadata and _metadata files. From the pyarrow docs:

Those files include information about the schema of the full dataset (for _common_metadata) and potentially all row group metadata of all files in the partitioned dataset as well (for _metadata). The actual files are metadata-only Parquet files. Note this is not a Parquet standard, but a convention set in practice by those frameworks.

This is especially helpful because you can read the _metadata file and then without any additional fetches you can do predicate pushdowns because you have the column statistics for every row group, and then you also have the exact byte ranges you want to fetch from. On that same page in the pyarrow docs they mention how to read and write.

On the Rust side, arrow2/parquet2 also have support for these _metadata files, see jorgecarleitao/parquet2#146 and jorgecarleitao/arrow2#1063.

[kylebarron]

I'm guessing that this would also likely be much faster once jorgecarleitao/parquet2#159 is implemented. Parquet stores its metadata in the footer of the file, and the last bytes of the file describe how many bytes the metadata take up. Right now arrow2/parquet2 require a HEAD request first to get the content length of the parquet file and then try to load the metadata. But this would be faster by:

1. Fetching the last N bytes of a Parquet file using a range request, where N is something like 64kb which often includes the entire metadata.
2. Checking the last 4 bytes of the buffer to get the length of the metadata
3. If N wasn't enough bytes, do one more fetch to get all the metadata.
4. Then parse the metadata

Maybe I'll find some time to try and make a PR for that in parquet2

[ritchie46]

@kylebarron I am working on a native polars parquet reader. Hope we can include this.

[winding-lines]

Current status:

• the scan_parquet code is DONE
• when collect-ing from python the physical_plan scan also needs to download the files, also DONE
• MISSING: an Tokio run-time needs to be started on the collect path, working on it

[winding-lines]

@ritchie46 this is ready for review. I am happy to change the API as per @talawahtech 's great feedback, if you agree with it. I am not yet familiar with the python API you desire for Polars :)

[ritchie46]

All those optimizations are default. We don't have to specify them.

[winding-lines]

Removed, thanks for letting me know.

[ritchie46]

We should not print the stdout.

[winding-lines]

Makes sense. The intent is to give some feedback for what is likely a slow operation.

I changed this to stderr, but if this is not acceptable I am happy to remove.

[ritchie46]

@kylebarron I am working on a native polars parquet reader. Hope we can include this.

[ritchie46]

I don't think ColumnStore is a good name. Can't local store whole row-groups?

[winding-lines]

@ritchie46 thanks for looking at the PR and I agree it's not a good name. To specifically answer you question, in local mode we don't actually store anything, it's just a pass through to the mmap-ed file.

MORE CONTEXT (sorry this is a bit long)

We are a bit in a tight spot:

• when running locally we want to use mmap
• when running remote we want to download only the parts of the file of interest (footer, metadata and then columns - this can be optimized a bit as per one of the comments)
• we currently want the async code to be as self-contained as possible

The ColumnStore is an abstraction that I introduced in an earlier PR to allow us to keep most of the code as-is and specially limit the spread of async. To achieve this this layer provides a level of indirection where the reader uses linear addressing, so the code stays as is.

For the local use case the ColumnStore is pass through:

• it maps the address requested by the reader to the actual address in the file
• the OS is doing its magic
• the code that you are looking at here is a no-op.

For the remote use case:

• we pre-download the regions of the file that we know will be accessed later on.
• This allows us to go in async mode for the period of the download and then come back to non-async mode.
• All the downloaded areas get store in a hashmap.
• When the reader executes and accesses an address we return the downloaded buffer.

I see a couple of ways forward:

1. allow the reader to run in async mode and actually do the download in the reader when the region of the file is needed
2. keep the current async island but come with a better name for the CloudStore

#1 is a bit of a departure from the current code architecture and you seem to prefer threads to async. Over the long run it may be inevitable though since we need to teach the planner the cost of operations.

If you want to go with #2 I could rename this to ColumMapper or CloudMapper, or some other name that resonates more with you.

Let me know what you think 🤗

[ritchie46]

Sorry for the delay (really busy). Great explanation. Could we add this as comments on top of the ColumnStore. Maybe we call it CloudMapper indeed.

I think we should first go for #2. I see some opportunities for some async in the streaming engine. But I don't want to overhaul our design and gradually experiment with this.

[winding-lines]

Thanks a lot Ritchie! Let's chain this PR on PR-6830. After we evaluate and possibly land PR-6830 we can come here and do a better job with the integration. This code will be much simpler if we have a Tokio runtime available when making it down to the download layer.

[stinodego]

I'm closing this pull request due to inactivity. Feel free to rebase and reopen and continue your work!

For reference, this is now possible in Python through fsspec. Though a native Rust implementation of Cloud integration would be welcome.