Spaces:
Sleeping
Sleeping
| """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 | |