core_lens.utils.polars_utils#

Polars scan helpers with predicate pushdown for entity materialisation.

Functions#

parquet_scan_path(→ str)

Return a glob-safe path for pl.scan_parquet.

cached_read_schema(→ Any)

Wrap pl.read_parquet_schema with caching.

collect_lf(→ polars.DataFrame)

Collect a LazyFrame using the best available backend.

scan_with_key_filter(→ polars.LazyFrame)

Return a pl.LazyFrame filtered to the given keys and optional time range.

Module Contents#

core_lens.utils.polars_utils.parquet_scan_path(path: str) → str#

Return a glob-safe path for pl.scan_parquet.

When path is a directory polars will reject it if the directory contains files with mixed extensions (e.g. a data_dictionary.csv sitting alongside Hive-partitioned .parquet files). This helper coerces a bare directory to <path>/**/*.parquet so only Parquet files are matched.

Parameters:

path (str) – Filesystem path or cloud URI to a Parquet file or directory.

Returns:

The original path if it already points to a file, otherwise <path>/**/*.parquet.

Return type:

str

core_lens.utils.polars_utils.cached_read_schema(path: str, storage_options: dict[str, Any] | None = None) → Any#

Wrap pl.read_parquet_schema with caching.

Avoids repeated Parquet footer reads for schema resolution in hot paths.

core_lens.utils.polars_utils.collect_lf(lf: polars.LazyFrame) → polars.DataFrame#

Collect a LazyFrame using the best available backend.

  • GPU present — executes via the RAPIDS cudf_polars streaming engine (GPUEngine(executor="streaming")). Handles datasets larger than VRAM through data partitioning.

  • No GPU — falls back to Polars’ built-in CPU streaming executor (collect(streaming=True)), which keeps memory usage low for large Parquet scans.

Use this function for all data scans (materialisation, geometry joins, similarity fetches). Tiny index scans that feed into subsequent in-process joins should stay as bare .collect() calls — the streaming path can occasionally change row ordering in ways that break those joins.

Parameters:

lf (pl.LazyFrame) – The lazy frame to collect.

Returns:

A materialised pl.DataFrame.

Return type:

pl.DataFrame

core_lens.utils.polars_utils.scan_with_key_filter(path: str, key_cols: list[str], key_values: polars.DataFrame | polars.LazyFrame, time_expr: polars.Expr | None = None, storage_options: dict[str, Any] | None = None) → polars.LazyFrame#

Return a pl.LazyFrame filtered to the given keys and optional time range.

Uses pl.scan_parquet with two predicate-pushdown layers:

  1. Key filter — restricts to entity instances whose key column(s) are in key_values. For a single-column key this is an is_in predicate pushed down to the Parquet reader. For composite keys each column is filtered independently (over-selects slightly, then pruned by the join at collect time).

  2. Time filter — an optional Polars expression appended with &, also pushed down if the Parquet file carries column statistics.

Parameters:

  • path (str) – Absolute path or cloud URI to a Parquet file.

  • key_cols (list[str]) – Column name(s) that form the entity’s unique key.

  • key_values (pl.DataFrame | pl.LazyFrame) – A narrow frame containing only the key column(s) with the exact values to retain.

  • time_expr (pl.Expr | None, optional) – An optional Polars filter expression for the time column, as produced by resolve_time_filter().

  • storage_options (dict[str, Any] | None, optional) – Cloud credential / configuration options forwarded to pl.scan_parquet. None uses ambient credentials.

Returns:

A pl.LazyFrame ready to be .collect()-ed.

Return type:

pl.LazyFrame