Skip to content

[Bug]: ColumnNotFoundError in scan_iceberg(reader_override="pyiceberg") on branching/self-join queries due to cached_ir projection overwrite #29345

Description

@Mrhs121

Checks

  • I have checked that this issue has not already been reported.
  • I have confirmed this bug exists on the latest version of Polars.

Reproducible example

Describe the bug

When executing a query where the same scan_iceberg(..., reader_override="pyiceberg") LazyFrame is branched into multiple paths with different column projections (e.g., in a semi-join or self-join where the left side requires ['id', 'http_method'] and the right side requires ['id']), query execution fails with ColumnNotFoundError.

However, the exact same query succeeds without issues when using reader_override="native".

Reproducible example

Here is a minimal, fully self-contained reproduction script using an in-memory SQLite catalog:

pip install "polars==1.44.2" "pyiceberg[sql-sqlite]" "pyarrow"

import tempfile
import shutil
from pathlib import Path
import polars as pl
from pyiceberg.catalog.sql import SqlCatalog
from pyiceberg.schema import Schema
from pyiceberg.types import LongType, NestedField, StringType

def run():
    tmp_dir = Path(tempfile.mkdtemp())
    try:
        # 1. Create a minimal local Iceberg table
        catalog = SqlCatalog(
            "default",
            uri=f"sqlite:///{tmp_dir}/cat.sqlite",
            warehouse=f"file://{tmp_dir}/wh"
        )
        catalog.create_namespace("default")
        schema = Schema(
            NestedField(1, "id", LongType()),
            NestedField(2, "http_method", StringType()),
            NestedField(3, "endpoint", StringType()),
        )
        table = catalog.create_table("default.test_table", schema=schema)
        pl.DataFrame({
            "id": [1, 2],
            "http_method": ["GET", "POST"],
            "endpoint": ["/api/v1", "/api/v2"],
        }).write_iceberg(table, mode="append")

        # 2. Scan with pyiceberg reader
        lf = pl.scan_iceberg(table, reader_override="pyiceberg")

        # 3. Branch into two paths with distinct projections:
        #    - Left branch needs: ['id', 'http_method']
        #    - Right branch needs: ['id', 'endpoint']
        left = lf.filter(pl.col("id") > 0)
        right = lf.filter(pl.col("endpoint").is_not_null()).select("id").limit(2)

        query = left.join(right, on="id", how="semi").select(["id", "http_method"])

        # Fails during collect
        query.collect()

    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)

if __name__ == "__main__":
    run()

Actual behavior and Traceback

Running the script above results in the following crash:

Traceback (most recent call last):
  File "demo.py", line 49, in <module>
    run()
  File "demo.py", line 43, in run
    query.collect()
  File "polars/_utils/deprecation.py", line 97, in wrapper
    return function(*args, **kwargs)
  File "polars/lazyframe/opt_flags.py", line 361, in wrapper
    return function(*args, **kwargs)
  File "polars/lazyframe/frame.py", line 2591, in collect
    return engine_.collect(
        self,
        engine=engine,
        streaming=streaming,
        post_opt_callback=post_opt_callback,
    )
  File "polars/lazyframe/engine.py", line 443, in collect
    return wrap_df(ldf.collect(self.name, callback))
polars.exceptions.ColumnNotFoundError: did not find column http_method, consider passing `missing_columns='insert'`

Log output

Traceback (most recent call last):
  File "demo.py", line 49, in <module>
    run()
  File "demo.py", line 43, in run
    query.collect()
  File "polars/_utils/deprecation.py", line 97, in wrapper
    return function(*args, **kwargs)
  File "polars/lazyframe/opt_flags.py", line 361, in wrapper
    return function(*args, **kwargs)
  File "polars/lazyframe/frame.py", line 2591, in collect
    return engine_.collect(
        self,
        engine=engine,
        streaming=streaming,
        post_opt_callback=post_opt_callback,
    )
  File "polars/lazyframe/engine.py", line 443, in collect
    return wrap_df(ldf.collect(self.name, callback))
polars.exceptions.ColumnNotFoundError: did not find column http_method, consider passing `missing_columns='insert'`

Root Cause

While PR #28468 (fixing #28465) resolved self-join panics by preventing duplicate AST node visits during IR traversal, its regression test only verified cases where both sides projected the exact same single column (['x']).

In queries where branching branches require heterogeneous projections:

  1. Projection Pushdown determines that the left branch requires ['id', 'http_method'] and the right branch only requires ['id'].
  2. In crates/polars-plan/src/plans/optimizer/expand_datasets.rs, FileScanIR::PythonDataset instances originating from the same LazyFrame share a single cached_ir: Arc<Mutex<Option<ExpandedDataset>>>.
  3. When the optimizer visits the left branch, it writes projection=['id', 'http_method'] into cached_ir.
  4. When the optimizer visits the right branch, it discovers a different projection and overwrites *cached_ir with the right branch's generator (which only yields ['id']).
  5. During physical plan lowering in crates/polars-stream/src/physical_plan/lower_ir.rs:823:
    FileScanIR::PythonDataset { dataset_object: _, cached_ir } => {
        let guard = cached_ir.lock().unwrap();
        let expanded_scan = guard.as_ref().unwrap().python_scan().unwrap();
        python_dataset_scan_to_reader_builder(expanded_scan)
    }
    The left branch's reader retrieves the overwritten generator (which only outputs ['id']).
  6. At runtime, ApplyExtraOps::initialize on the left branch verifies the DataFrame emitted by the generator against its expected schema, notices that http_method is missing, and raises ColumnNotFoundError.

Expected behavior

The query should complete successfully and return the filtered rows with ['id', 'http_method'], identically to how reader_override="native" behaves.

Installed versions

  • Polars Version: 1.44.2 (and verified present on latest main commit 3127ec9c82)
  • Python Version: 3.13
  • PyIceberg Version: 0.12.0
  • PyArrow Version: 19.x / 25.x
  • Platform: macOS (arm64) / Linux

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    A-io-icebergRelated to Apache Iceberg tables.bugSomething isn't workingneeds triageAwaiting prioritization by a maintainerpythonRelated to Python Polars

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions