[Python] Fastest way to join a small database and a big database in Python

Suppose I have two dataframes, small_df and huge_df in two parquet files. small_df has about 500K rows and huge_df has more than 10 billion lines. They both have a column "A", and small_df is not sorted on "A", but huge_df is. I want to sort small_df and left-join it with huge_df on "A" in the fastest possible way. I have a lot of memory available and a lot of CPUs (many hundreds) so resources are not a problem, I just want to optimize for time. How can I do it in Python?

Here is a mwe:

import polars as pl
import numpy as np
import time
import os

# Print parquet file sizes
def format_bytes(bytes_size):
"""Convert bytes to human readable format"""
for unit in ['B', 'KB', 'MB', 'GB']:
if bytes_size < 1024.0:
return f"{bytes_size:.1f} {unit}"
bytes_size /= 1024.0
return f"{bytes_size:.1f} TB"

def create_test_data(small_size=500_000, huge_size=200_000_000, temp_dir="./temp_data"):
"""Create sample data and save to parquet files for realistic I/O testing"""
os.makedirs(temp_dir, exist_ok=True)

# Create huge_df (already sorted on column 'A') with 10 columns
huge_df = pl.DataFrame({
'A': range(1, huge_size + 1), # Sequential = sorted
'huge_value': np.random.randn(huge_size),
'huge_category': np.random.choice(['X', 'Y', 'Z'], huge_size),
'huge_score': np.random.uniform(0, 100, huge_size),
'huge_flag': np.random.choice([True, False], huge_size),
'huge_price': np.random.exponential(50, huge_size),
'huge_rating': np.random.choice([1, 2, 3, 4, 5], huge_size),
'huge_region': np.random.choice(['North', 'South', 'East', 'West'], huge_size),
'huge_count': np.random.poisson(10, huge_size),
'huge_percentage': np.random.beta(2, 5, huge_size) * 100
})

# Create small_df (unsorted on column 'A', with 80% overlap)
overlap_size = int(small_size * 0.8)
small_keys = (
list(np.random.choice(huge_size, overlap_size, replace=False)) + # Keys that exist
list(range(huge_size + 1, huge_size + small_size - overlap_size + 1)) # Keys that don't
)
np.random.shuffle(small_keys) # Make it unsorted

# Create small_df (unsorted) with 20 columns
small_df = pl.DataFrame({
'A': small_keys,
'small_value': np.random.randn(small_size),
'small_flag': np.random.choice([True, False], small_size),
'small_category': np.random.choice(['Alpha', 'Beta', 'Gamma', 'Delta'], small_size),
'small_score': np.random.uniform(0, 1000, small_size),
'small_price': np.random.lognormal(4, 1, small_size),
'small_quantity': np.random.randint(1, 100, small_size),
'small_status': np.random.choice(['Active', 'Inactive', 'Pending'], small_size),
'small_weight': np.random.gamma(2, 2, small_size),
'small_temperature': np.random.normal(20, 5, small_size),
'small_pressure': np.random.exponential(1.5, small_size),
'small_density': np.random.uniform(0.8, 1.2, small_size),
'small_ph': np.random.normal(7, 0.5, small_size),
'small_conductivity': np.random.exponential(10, small_size),
'small_viscosity': np.random.gamma(1.5, 2, small_size),
'small_opacity': np.random.beta(2, 3, small_size),
'small_hardness': np.random.choice([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], small_size),
'small_color': np.random.choice(['Red', 'Blue', 'Green', 'Yellow', 'Orange'], small_size),
'small_texture': np.random.choice(['Smooth', 'Rough', 'Grainy', 'Silky'], small_size),
'small_timestamp': np.random.randint(1600000000, 1700000000, small_size)
})

small_path = f"{temp_dir}/small_df.parquet"
huge_path = f"{temp_dir}/huge_df.parquet"

small_df.write_parquet(small_path)
huge_df.write_parquet(huge_path)

small_size_bytes = os.path.getsize(small_path)
huge_size_bytes = os.path.getsize(huge_path)

print(f"Small parquet file: {format_bytes(small_size_bytes)}")
print(f"Huge parquet file: {format_bytes(huge_size_bytes)}")

return small_path, huge_path

def naive_approach(small_parquet_path, huge_parquet_path):
start_time = time.time()

# Scan parquet files lazily - no data loaded into memory yet
small_lf = pl.scan_parquet(small_parquet_path)
huge_lf = pl.scan_parquet(huge_parquet_path)

# Build query plan - Polars will optimize the entire pipeline
result_lf = (
small_lf
.sort("A") # Sort small_df on column A
.join(huge_lf, on="A", how="left") # Left join with huge_df
)

# Execute optimized plan (includes I/O time)
result_df = result_lf.collect()

print(f"Join completed in {time.time() - start_time:.2f}s")
return result_df

def main():
print("Creating test data...")
small_path, huge_path = create_test_data()

# Quick check of file sizes
small_df_info = pl.scan_parquet(small_path).select(pl.len()).collect().item()
huge_df_info = pl.scan_parquet(huge_path).select(pl.len()).collect().item()

print(f"Small parquet: {small_df_info:,} rows (unsorted)")
print(f"Huge parquet: {huge_df_info:,} rows (sorted)")

print("\n--- Polars Sort and Join ---")
result = naive_approach(small_path, huge_path)

print(f"Result shape: {result.shape}")
print(f"Successful joins: {result.filter(pl.col('huge_value').is_not_null()).height}")

# Cleanup
try:
os.remove(small_path)
os.remove(huge_path)
os.rmdir(os.path.dirname(small_path))
print("Cleaned up temporary files")
except OSError:
print(f"Note: Temporary files remain at {os.path.dirname(small_path)}")

if __name__ == "__main__":
# Run: pip install polars numpy
main()


On my Thinkpad laptop, so not in production, it gives these results:

Creating test data...
Small parquet file: 44.8 MB
Huge parquet file: 6.1 GB
Small parquet: 500,000 rows (unsorted)
Huge parquet: 200,000,000 rows (sorted)

--- Polars Sort and Join ---
Join completed in 12.71s
Result shape: (500000, 29)
Successful joins: 400000
Cleaned up temporary files

Continue reading...