Spaces:
Sleeping
Sleeping
File size: 1,736 Bytes
c21180c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | """CSV loading: DuckDB registration + Polars sampling."""
from __future__ import annotations
from pathlib import Path
from typing import Any
import duckdb
import polars as pl
def register_csv(con: duckdb.DuckDBPyConnection, table_name: str, csv_path: str) -> dict[str, Any]:
"""Register a CSV file as a DuckDB table and return metadata."""
con.execute(
f"CREATE OR REPLACE TABLE \"{table_name}\" AS SELECT * FROM read_csv_auto('{csv_path}', header=true, ignore_errors=true)"
)
row_count = con.execute(f'SELECT COUNT(*) FROM "{table_name}"').fetchone()[0]
col_info = con.execute(
f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}'"
).fetchall()
return {
"table_name": table_name,
"row_count": row_count,
"columns": [{"name": c[0], "dtype": c[1]} for c in col_info],
}
def read_csv_sample(path: str, n: int = 10_000) -> pl.DataFrame:
"""Read a CSV sample using Polars for fast preview."""
return pl.read_csv(path, n_rows=n, ignore_errors=True)
def read_csv_lazy(path: str) -> pl.LazyFrame:
"""Read a CSV as a Polars LazyFrame for efficient transformations."""
return pl.scan_csv(path, ignore_errors=True)
def export_to_csv(con: duckdb.DuckDBPyConnection, table_name: str, output_path: str) -> str:
"""Export a DuckDB table to CSV."""
con.execute(f"COPY \"{table_name}\" TO '{output_path}' (HEADER, DELIMITER ',')")
return output_path
def export_to_parquet(con: duckdb.DuckDBPyConnection, table_name: str, output_path: str) -> str:
"""Export a DuckDB table to Parquet."""
con.execute(f"COPY \"{table_name}\" TO '{output_path}' (FORMAT PARQUET)")
return output_path
|