从 Pandas 到 Polars 四十六：使用Polars读取和写入S3数据

在本文中，我们将看到如何使用Polars从S3中的CSV或Parquet文件读取和写入数据。同时，我们还将了解如何在下载前对S3上的文件进行过滤，以减少跨网络传输的数据量。

写入文件到S3

我们将创建一个包含3列的简单DataFrame。我们将使用s3fs库将其写入S3中的CSV和Parquet文件。s3fs库允许您以类似于在本地文件系统上工作的语法来读取和写入S3中的文件。

bucket_name = "my_bucket"
csv_key = "test_write.csv"
parquet_key = "test_write.parquet"
fs = s3fs.S3FileSystem()  
df = pl.DataFrame(
    {
        "foo": [1, 2, 3, 4, 5],
        "bar": [6, 7, 8, 9, 10],
        "ham": ["a", "b", "c", "d", "e"],
    }
)
with fs.open(f"{bucket_name}/{csv_key}", mode="wb") as f:
    df.write_csv(f)
with fs.open(f"{bucket_name}/{parquet_key}", mode="wb") as f:
    df.write_parquet(f)

如果你可以选择的话，我推荐使用Parquet格式，因为它具有更小的文件大小，可以保留数据类型（dtypes），并且使后续读取更快。

从S3读取文件

我们可以使用Polars的pl.read_csv函数从S3中读取文件。

df_csv = pl.read_csv(f"s3://{bucket}/{csv_key}")
df_parquet = pl.read_parquet(f"s3://{bucket}/{parquet_key}")

Polars内部使用ffspec将远程文件读取到内存缓冲区中，然后将缓冲区中的数据读入DataFrame。这是一种快速的方法，但它确实意味着整个文件都被读入内存。对于小文件来说这没问题，但对于大文件来说可能会很慢并且占用大量内存。