Skip to content

Commit

Permalink
c
Browse files Browse the repository at this point in the history
  • Loading branch information
nameexhaustion committed Oct 21, 2024
1 parent f3eba22 commit 6b08139
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
15 changes: 15 additions & 0 deletions crates/polars-plan/src/plans/optimizer/projection_pushdown/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -511,6 +511,21 @@ impl ProjectionPushDown {
file_options.row_index = None;
}
};

if let Some(col_name) = &file_options.include_file_paths {
if output_schema
.as_ref()
.map_or(false, |schema| !schema.contains(col_name))
{
// Need to remove it from the input schema so
// that projection indices are correct.
let mut file_schema = Arc::unwrap_or_clone(file_info.schema);
file_schema.shift_remove(col_name);
file_info.schema = Arc::new(file_schema);
file_options.include_file_paths = None;
}
};

let lp = Scan {
sources,
file_info,
Expand Down
8 changes: 8 additions & 0 deletions py-polars/tests/unit/io/test_scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -801,3 +801,11 @@ def test_scan_double_collect_row_index_invalidates_cached_ir_18892() -> None:
schema={"index": pl.UInt32, "a": pl.Int64},
),
)


def test_scan_include_file_paths_respects_projection_pushdown() -> None:
q = pl.scan_csv("a,b,c\na1,b1,c1".encode(), include_file_paths="path_name").select(
["a", "b"]
)

assert_frame_equal(q.collect(), pl.DataFrame({"a": "a1", "b": "b1"}))

0 comments on commit 6b08139

Please sign in to comment.