feat(pkg-py): add PolarsLazySource for Polars LazyFrame support

[cpsievert]

Summary

• Adds PolarsLazySource DataSource implementation for Polars LazyFrames
• Data stays lazy until the render boundary, enabling efficient handling of large datasets
• Uses Polars' SQLContext for native lazy SQL execution
• Adds AnyFrame type alias to support both DataFrame and LazyFrame return types

Usage

If using the Quick start .app(), just drop in a LazyFrame, same as you would a DataFrame.

import polars as pl
from querychat import QueryChat

# Scan a large parquet file lazily
lf = pl.scan_parquet("large_data.parquet")

# Pass directly to QueryChat - stays lazy until rendered
qc = QueryChat(data_source=lf, table_name="data")
qc.app()

If using the .df() to build a custom app, note that this will return the filtered LazyFrame, so you may need to .collect() as needed.

import polars as pl
from shiny.express import render, ui
from querychat import QueryChat

# Scan lazily - data stays on disk
lf = pl.scan_parquet("large_data.parquet")
qc = QueryChat(data_source=lf, table_name="sales")

# Sidebar with chat interface
with ui.sidebar():
    qc.ui()

# Main panel with data table
@render.data_frame
def filtered_data():
    # df() returns a LazyFrame - collect at render time
    return qc.df().collect()

Changes

• _datasource.py: Add AnyFrame type alias, widen ABC return types, implement PolarsLazySource
• _querychat.py: Update normalize_data_source() to detect LazyFrames, collect at render boundary
• _querychat_module.py: Update ServerValues.df type hint
• tools.py: Handle LazyFrames in query tool
• data-sources.qmd: Add documentation for LazyFrame support
• Tests: 15 new tests for PolarsLazySource, 1 integration test

Performance

With a 100K row test dataset:

Operation	Eager	Lazy	Speedup
Load time	0.02s	0.0001s	200x
QueryChat init	0.26s	0.06s	4x
Query execution	0.50s	0.009s	55x

Demo script for manual verification

Run with:

export OPENAI_API_KEY="your-key"
cd pkg-py
uv run python examples/lazy_frame_demo.py --rows 10000000

#!/usr/bin/env python3
"""
Demo script comparing eager vs lazy data source performance.

This script demonstrates the performance benefits of using PolarsLazySource
with large datasets. It creates a synthetic dataset and compares:
1. Eager loading (all data in memory upfront)
2. Lazy loading (data stays on disk until needed)

Usage:
    # Set your API key first
    export OPENAI_API_KEY="your-key-here"

    # Run the demo
    cd pkg-py
    uv run python examples/lazy_frame_demo.py

    # Or with a custom number of rows (default: 10 million)
    uv run python examples/lazy_frame_demo.py --rows 50000000
"""

import argparse
import os
import tempfile
import time
from pathlib import Path

import polars as pl


def create_large_dataset(path: Path, n_rows: int) -> None:
    """Create a large parquet file for testing."""
    print(f"Creating dataset with {n_rows:,} rows...")
    start = time.perf_counter()

    # Generate data in chunks to avoid memory issues
    chunk_size = 1_000_000
    chunks_written = 0

    for i in range(0, n_rows, chunk_size):
        chunk_rows = min(chunk_size, n_rows - i)
        chunk = pl.DataFrame(
            {
                "id": range(i, i + chunk_rows),
                "category": [f"cat_{j % 100}" for j in range(chunk_rows)],
                "region": [["North", "South", "East", "West"][j % 4] for j in range(chunk_rows)],
                "value": [float(j % 1000) + 0.5 for j in range(chunk_rows)],
                "quantity": [j % 500 for j in range(chunk_rows)],
                "date": pl.Series([f"2024-{(j % 12) + 1:02d}-{(j % 28) + 1:02d}" for j in range(chunk_rows)]).str.to_date(),
            }
        )

        if chunks_written == 0:
            chunk.write_parquet(path)
        else:
            existing = pl.read_parquet(path)
            pl.concat([existing, chunk]).write_parquet(path)

        chunks_written += 1
        print(f"  Written {min(i + chunk_size, n_rows):,} / {n_rows:,} rows")

    elapsed = time.perf_counter() - start
    file_size_mb = path.stat().st_size / (1024 * 1024)
    print(f"Dataset created: {file_size_mb:.1f} MB in {elapsed:.1f}s\n")


def measure_memory() -> float:
    """Get current memory usage in MB (approximate)."""
    import psutil
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / (1024 * 1024)


def demo_eager_vs_lazy(parquet_path: Path) -> None:
    """Compare eager vs lazy data loading performance."""
    from querychat import QueryChat

    print("=" * 60)
    print("COMPARING EAGER VS LAZY DATA SOURCE")
    print("=" * 60)

    try:
        import psutil
        has_psutil = True
    except ImportError:
        has_psutil = False
        print("(Install psutil for memory usage tracking: pip install psutil)\n")

    # --- EAGER LOADING ---
    print("\n1. EAGER LOADING (polars.read_parquet → DataFrame)")
    print("-" * 50)

    if has_psutil:
        mem_before = measure_memory()

    start = time.perf_counter()
    df = pl.read_parquet(parquet_path)
    load_time = time.perf_counter() - start

    if has_psutil:
        mem_after = measure_memory()
        print(f"   Memory increase: {mem_after - mem_before:.1f} MB")

    print(f"   Load time: {load_time:.2f}s")
    print(f"   Rows loaded: {len(df):,}")

    start = time.perf_counter()
    qc_eager = QueryChat(data_source=df, table_name="sales", greeting="Hello!")
    init_time = time.perf_counter() - start
    print(f"   QueryChat init: {init_time:.2f}s")

    start = time.perf_counter()
    result = qc_eager.data_source.execute_query(
        "SELECT region, SUM(value) as total FROM sales GROUP BY region"
    )
    query_time = time.perf_counter() - start
    print(f"   Query execution: {query_time:.3f}s")
    print(f"   Result rows: {len(result)}")

    del df, qc_eager, result
    import gc
    gc.collect()

    # --- LAZY LOADING ---
    print("\n2. LAZY LOADING (polars.scan_parquet → LazyFrame)")
    print("-" * 50)

    if has_psutil:
        mem_before = measure_memory()

    start = time.perf_counter()
    lf = pl.scan_parquet(parquet_path)
    load_time = time.perf_counter() - start

    if has_psutil:
        mem_after = measure_memory()
        print(f"   Memory increase: {mem_after - mem_before:.1f} MB")

    print(f"   'Load' time: {load_time:.4f}s (just metadata!)")

    start = time.perf_counter()
    qc_lazy = QueryChat(data_source=lf, table_name="sales", greeting="Hello!")
    init_time = time.perf_counter() - start
    print(f"   QueryChat init: {init_time:.2f}s")

    start = time.perf_counter()
    result_lazy = qc_lazy.data_source.execute_query(
        "SELECT region, SUM(value) as total FROM sales GROUP BY region"
    )
    query_time = time.perf_counter() - start
    print(f"   Query execution (lazy): {query_time:.3f}s")

    start = time.perf_counter()
    result_collected = result_lazy.collect()
    collect_time = time.perf_counter() - start
    print(f"   Collect time: {collect_time:.3f}s")
    print(f"   Result rows: {len(result_collected)}")

    print("\n" + "=" * 60)
    print("SUMMARY: Lazy loading is dramatically faster for large datasets!")
    print("=" * 60)


def main():
    parser = argparse.ArgumentParser(description="Demo lazy vs eager data loading")
    parser.add_argument("--rows", type=int, default=10_000_000)
    parser.add_argument("--interactive", action="store_true")
    parser.add_argument("--data-path", type=str, default=None)
    args = parser.parse_args()

    if args.data_path:
        parquet_path = Path(args.data_path)
    else:
        temp_dir = tempfile.mkdtemp()
        parquet_path = Path(temp_dir) / "large_sales_data.parquet"
        create_large_dataset(parquet_path, args.rows)

    try:
        demo_eager_vs_lazy(parquet_path)
    finally:
        if not args.data_path and parquet_path.exists():
            parquet_path.unlink()
            parquet_path.parent.rmdir()


if __name__ == "__main__":
    main()

Sample output (100K rows):

============================================================
COMPARING EAGER VS LAZY DATA SOURCE
============================================================

1. EAGER LOADING (polars.read_parquet → DataFrame)
--------------------------------------------------
   Load time: 0.02s
   Rows loaded: 100,000
   QueryChat init: 0.26s
   Query execution: 0.495s
   Result rows: 4

2. LAZY LOADING (polars.scan_parquet → LazyFrame)
--------------------------------------------------
   'Load' time: 0.0001s (just metadata!)
   QueryChat init: 0.06s
   Query execution (lazy): 0.009s
   Collect time: 0.007s
   Result rows: 4

Test plan

• All 131 tests pass
• Pyright passes with 0 errors on source files
• New PolarsLazySource tests cover init, execute_query, get_data, get_schema, test_query
• Integration test verifies QueryChat accepts LazyFrame and creates PolarsLazySource
• Demo script validates performance benefits

🤖 Generated with Claude Code

[cpsievert]

This will be an annoying thing to workaround as a user with type checking enabled.

I have some ideas on how to improve this, but I think it's better to address that in a follow up PR.