DuckDB sampling from parquet and hive partitioned parquet files

duckdb_sampling.ipynb

duckdb_sampling.py

import duckdb
import numpy as np
import pandas as pd
from tempfile import TemporaryDirectory, NamedTemporaryFile

conn = duckdb.connect(":memory:")
conn.execute("SET threads = 1;")

# generate data
np.random.seed(42)
size = 1_000_000
df = pd.DataFrame(
    {
        "range": range(size),
        "bin": np.random.randint(0, 10, size=size),
    }
)

print("sampling from df")
print(
    conn.sql("""
  SELECT range
  FROM df
  USING SAMPLE reservoir(10 ROWS)
  REPEATABLE(42);
  """).fetchall()
)

# write to single parquet
parquet_file = NamedTemporaryFile(suffix=".parquet")
conn.sql(f"""
  COPY (FROM df)
  TO '{parquet_file.name}'
  (FORMAT 'parquet', OVERWRITE);
""")

print("sampling from single parquet")
print(
    conn.sql(f"""
  SELECT range
  FROM '{parquet_file.name}'
  USING SAMPLE reservoir(10 ROWS)
  REPEATABLE(42);
  """).fetchall()
)

parquet_file.close()

# write a hive partitioned parquet files
hive_path = TemporaryDirectory()
conn.sql(f"""
  COPY (FROM df)
  TO '{hive_path.name}'
  (FORMAT 'parquet', PARTITION_BY bin, OVERWRITE);
  """)

print("sampling from hive partitioned parquet files")
print(
    conn.sql(f"""
  SELECT range
  FROM read_parquet('{hive_path.name}/*/*.parquet', hive_partitioning = true)
  USING SAMPLE reservoir(10 ROWS)
  REPEATABLE(42);
  """).fetchall()
)

hive_path.cleanup()