Spaces:
Sleeping
Sleeping
| """DataFrameRef: thin wrapper over DuckDB table or Polars LazyFrame.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field | |
| from typing import Any, Literal | |
| import duckdb | |
| import polars as pl | |
| class DataFrameRef: | |
| """Reference to a DataFrame backed by either DuckDB or Polars.""" | |
| backend: Literal["duckdb", "polars"] | |
| name: str | None = None # DuckDB table name | |
| lf: pl.LazyFrame | None = None # Polars LazyFrame | |
| con: duckdb.DuckDBPyConnection | None = None # DuckDB connection (for duckdb backend) | |
| source_path: str | None = None # Original CSV path | |
| metadata: dict[str, Any] = field(default_factory=dict) | |
| def head(self, n: int = 50, offset: int = 0) -> pl.DataFrame: | |
| """Return first n rows as a Polars DataFrame.""" | |
| if self.backend == "duckdb" and self.name and self.con: | |
| result = self.con.execute( | |
| f'SELECT * FROM "{self.name}" LIMIT {n} OFFSET {offset}' | |
| ).fetchall() | |
| col_names = [d[0] for d in self.con.execute(f'SELECT * FROM "{self.name}" LIMIT 0').description] | |
| return pl.DataFrame(result, schema=col_names, orient="row") | |
| elif self.backend == "polars" and self.lf is not None: | |
| return self.lf.slice(offset, n).collect() | |
| raise ValueError("Invalid DataFrameRef: no backend available") | |
| def schema(self) -> list[dict[str, str]]: | |
| """Return column names and types.""" | |
| if self.backend == "duckdb" and self.name and self.con: | |
| rows = self.con.execute( | |
| f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{self.name}'" | |
| ).fetchall() | |
| return [{"name": r[0], "dtype": r[1]} for r in rows] | |
| elif self.backend == "polars" and self.lf is not None: | |
| schema = self.lf.schema | |
| return [{"name": name, "dtype": str(dtype)} for name, dtype in schema.items()] | |
| raise ValueError("Invalid DataFrameRef") | |
| def to_polars_df(self) -> pl.DataFrame: | |
| """Convert to a full Polars DataFrame.""" | |
| if self.backend == "duckdb" and self.name and self.con: | |
| result = self.con.execute(f'SELECT * FROM "{self.name}"').fetchall() | |
| col_names = [d[0] for d in self.con.execute(f'SELECT * FROM "{self.name}" LIMIT 0').description] | |
| return pl.DataFrame(result, schema=col_names, orient="row") | |
| elif self.backend == "polars" and self.lf is not None: | |
| return self.lf.collect() | |
| raise ValueError("Invalid DataFrameRef") | |
| def to_lazy(self) -> pl.LazyFrame: | |
| """Get as Polars LazyFrame.""" | |
| if self.backend == "polars" and self.lf is not None: | |
| return self.lf | |
| elif self.backend == "duckdb" and self.name and self.con: | |
| df = self.to_polars_df() | |
| return df.lazy() | |
| raise ValueError("Invalid DataFrameRef") | |
| def row_count(self) -> int: | |
| """Return total row count.""" | |
| if self.backend == "duckdb" and self.name and self.con: | |
| return self.con.execute(f'SELECT COUNT(*) FROM "{self.name}"').fetchone()[0] | |
| elif self.backend == "polars" and self.lf is not None: | |
| return self.lf.select(pl.len()).collect().item() | |
| return 0 | |
| def materialize_to_duckdb(self, table_name: str, con: duckdb.DuckDBPyConnection) -> DataFrameRef: | |
| """Materialize a Polars LazyFrame into a DuckDB table.""" | |
| if self.backend == "polars" and self.lf is not None: | |
| df = self.lf.collect() | |
| # Use Apache Arrow for efficient transfer | |
| arrow_tbl = df.to_arrow() | |
| con.register("pl_df", arrow_tbl) | |
| con.execute(f'CREATE OR REPLACE TABLE "{table_name}" AS SELECT * FROM pl_df') | |
| con.unregister("pl_df") | |
| return DataFrameRef(backend="duckdb", name=table_name, con=con, source_path=self.source_path) | |
| return self | |
| def clone(self) -> DataFrameRef: | |
| """Create a shallow clone.""" | |
| return DataFrameRef( | |
| backend=self.backend, | |
| name=self.name, | |
| lf=self.lf, | |
| con=self.con, | |
| source_path=self.source_path, | |
| metadata=dict(self.metadata), | |
| ) | |