Silent row drop in db.sql(query).pl(lazy=True) when the LazyFrame is self-rejoined and collected via engine="streaming" with a window expression #452

"""Minimal reproducer for a silent ~80% row drop in duckdb-python's
polars bridge under collect(engine="streaming").

When two ``db.sql(query).pl(lazy=True)`` LazyFrames are joined, the
joined LazyFrame is referenced twice in the plan (a self-rejoin), a
window expression with ``.over("g")`` sits downstream of an explicit
``.sort(["g","t"])``, and the plan is collected via
``collect(engine="streaming")``, the output drops ~80% of rows compared
to the pure ``pl.scan_parquet`` reference. The sort guarantees a stable
within-group order so the ``rolling_sum`` is deterministic — but the
row count is still ~80% short of the reference.

Example output on the repro below (50,597,009 rows total in reference):

    A) pl.scan_parquet, in-memory engine
       shape=(50597009, 3)  y non-null=45,647,009

    B) db.sql(...).pl(lazy=True), streaming engine, .sort(['g','t']) added
       shape=(10400465, 3)  y non-null=9,377,597   shape_match=False

The same plan against ``pl.scan_parquet`` LazyFrames produces the full
50.6M rows. The duckdb bridge is dropping rows during streaming.

Tested with:
    polars  1.40.1
    duckdb  1.5.1 / 1.5.2 / 1.6.0.dev12 (all reproduce)
    pyarrow latest

Run:
    uv run --with 'polars==1.40.1' --with numpy --with pyarrow --with duckdb \\
        python duckdb_polars_streaming_repro.py
"""
from __future__ import annotations

import shutil
import tempfile
from pathlib import Path

import duckdb
import numpy as np
import polars as pl

print(f"polars {pl.__version__}  duckdb {duckdb.__version__}")

# ---------------------------------------------------------------------------
# 1) Two parquet files with variable-length groups.
#    `left`:  (g, t, x)  — group id, index, payload
#    `right`: (g, t)     — same key columns, no payload
# Variable per-group length matters: uniform-length groups don't reproduce
# at this scale.
# ---------------------------------------------------------------------------
N_GROUPS = 50_000
rng = np.random.default_rng(42)
group_lens = np.clip(
    rng.lognormal(mean=6.8, sigma=0.5, size=N_GROUPS).astype(int), 30, 2900
)
N = int(group_lens.sum())
print(f"rows: {N:,}  groups: {N_GROUPS:,}")

g = np.repeat(np.arange(N_GROUPS, dtype=np.int32), group_lens)
t = np.concatenate([np.arange(n, dtype=np.int32) for n in group_lens])
x = rng.uniform(-1.0, 1.0, N).astype(np.float32)

tmpdir = Path(tempfile.mkdtemp(prefix="repro_"))
left_path = tmpdir / "left.parquet"
right_path = tmpdir / "right.parquet"
pl.DataFrame({"g": g, "t": t, "x": x}).write_parquet(
    left_path, row_group_size=200_000
)
pl.DataFrame({"g": g, "t": t}).write_parquet(
    right_path, row_group_size=200_000
)
del g, t, x


# ---------------------------------------------------------------------------
# 2) The plan:
#    - `left` and `right` are joined into `joined`.
#    - `joined` is referenced TWICE: once directly, once via a derived
#      LazyFrame that selects the group keys uniquely.
#    - A rolling-sum window over `g` is applied at the end.
# ---------------------------------------------------------------------------
def build(left_lf: pl.LazyFrame, right_lf: pl.LazyFrame) -> pl.LazyFrame:
    joined = left_lf.join(right_lf, on=["g", "t"], how="inner")
    keys = joined.select("g").unique()                # ← reference #1
    plan = joined.join(keys, on="g")                  # ← reference #2
    # Explicit sort by (partition key, within-partition key) so the
    # downstream rolling_sum().over("g") sees within-g rows in t-order.
    # Inner joins do NOT preserve sort order, so this sort is required
    # for the .over result to be deterministic at all.
    return plan.sort(["g", "t"]).select(
        "g",
        "t",
        pl.col("x").rolling_sum(window_size=100).over("g").alias("y"),
    )


# ---------------------------------------------------------------------------
# 3) Reference: pure pl.scan_parquet, in-memory engine.
# ---------------------------------------------------------------------------
print("\nA) pl.scan_parquet, in-memory engine")
ref = build(pl.scan_parquet(left_path), pl.scan_parquet(right_path)).collect()
ref_shape = ref.shape
ref_nn = ref["y"].drop_nulls().drop_nans().len()
print(f"   shape={ref_shape}  y non-null={ref_nn:,}")

# ---------------------------------------------------------------------------
# 4) Bug path: db.sql(...).pl(lazy=True) + collect(engine="streaming").
# Single run — the sort makes the output deterministic across runs.
# ---------------------------------------------------------------------------
print("\nB) db.sql(...).pl(lazy=True), streaming engine, sorted plan")
db_l = duckdb.connect(":memory:")
db_r = duckdb.connect(":memory:")
try:
    left_lf = db_l.sql(
        f"select * from read_parquet('{left_path}')"
    ).pl(lazy=True)
    right_lf = db_r.sql(
        f"select * from read_parquet('{right_path}')"
    ).pl(lazy=True)
    out = build(left_lf, right_lf).collect(engine="streaming")
    nn = out["y"].drop_nulls().drop_nans().len()
    print(
        f"   shape={out.shape}  y non-null={nn:,}  "
        f"shape_match={out.shape == ref_shape}  nn_match={nn == ref_nn}"
    )
finally:
    db_l.close()
    db_r.close()

shutil.rmtree(tmpdir)