CODEWITHSUNDEEP
pythonpandasparquetpyarrow
PaleNeutron
PaleNeutron

Reputation: 3235

How to read a part of parquet dataset into pandas?

I have a huge dataframe and want to split it into small files for better performance. Here is the example code to write. BUT I can not just read a small pieces from it without loading whole dataframe into memory.

import pandas as pd
import os

# Create a sample DataFrame with daily frequency
data = {
    "timestamp": pd.date_range(start="2023-01-01", periods=1000, freq="D"),
    "value": range(100)
}
df = pd.DataFrame(data)

# Add a column for year (to use as a partition key)
df["year"] = df["timestamp"].dt.year
df["month"] = df["timestamp"].dt.month

# Use the join method to expand the DataFrame (Cartesian product with a multiplier)
multiplier = pd.DataFrame({"replica": range(100)})  # Create a multiplier DataFrame
expanded_df = df.join(multiplier, how="cross")  # Cartesian product using cross join

# Define the output directory
output_dir = "output_parquet"

# Save the expanded DataFrame to Parquet with year-based partitioning
expanded_df.to_parquet(
    output_dir, 
    partition_cols=["year", "month"],  # Specify the partition column
)

Which is the best way to read from the dataset if I only need data from 2023-12-01 to 2024-01-31?

Upvotes: 1

Views: 78

Answers (1)

Parman M. Alizadeh
Parman M. Alizadeh

Reputation: 1553

You can load your dataset selectively if your parquet output is properly partitioned. You can use libraries like PyArrow to let you filter the data at the file or partition level to make sure only the relevant data is loaded in to memory.

Here's how you can do it using pyarrow.dataset:

import pyarrow.dataset as ds

dataset = ds.dataset("output_parquet", format="parquet", partitioning="hive")
filtered_table = dataset.to_table(filter=(ds.field("year") == 2023) & (ds.field("month") == 12)) # Or any other desired condition
filtered_df = filtered_table.to_pandas()

You may also use pyarrow.parquet.ParquetDataset to achieve the same goal, but it's less optimized and a bit outdated:

import pyarrow.parquet as pq

dataset = pq.ParquetDataset("output_parquet", filters=[("year", "=", 2023), ("month", "=", 12)])
table = dataset.read()
df = table.to_pandas()

Upvotes: 3

Related Questions