diff --git a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/hive.rs b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/hive.rs index 7d8308825961..9ab5b9d360f2 100644 --- a/crates/polars-plan/src/plans/optimizer/predicate_pushdown/hive.rs +++ b/crates/polars-plan/src/plans/optimizer/predicate_pushdown/hive.rs @@ -44,6 +44,16 @@ fn get_partitions(hive_df: &DataFrame) -> Vec { split_df_as_ref(hive_df, std::cmp::min(n_parts, hive_df.height()), false) } +// We must deduplicate so that we get unique files (keys) per partition. +#[cfg(feature = "is_in")] +fn unique_key_frame(hive_df: &DataFrame, key_name: &PlSmallStr) -> PolarsResult { + hive_df + .column(key_name)? + .clone() + .into_frame() + .unique_stable(None, UniqueKeepStrategy::First, None) +} + #[allow(clippy::too_many_arguments)] pub fn rewrite_hive( ir: IR, @@ -78,20 +88,22 @@ pub fn rewrite_hive( // We do that by pushing down an is_in predicate // Later in the optimizer we prune the hive paths // based on all the predicates. - let mut hive_col = None; + // Any hive column is valid to split on, prefer the leftmost as it is the coarsest. + let mut hive_col: Option<(usize, PlSmallStr)> = None; let hive_schema = hive.schema(); for e in keys.iter() { let key = expr_arena.get(e.node()); if let AExpr::Column(name) = key { - if hive_schema.index_of(name) == Some(0) { - hive_col = Some(name.clone()); - break; + if let Some(idx) = hive_schema.index_of(name) + && hive_col.as_ref().is_none_or(|(best, _)| idx < *best) + { + hive_col = Some((idx, name.clone())); } } } - if let Some(key_name) = hive_col { - let hive_df = hive.df().select_at_idx(0).unwrap().clone().into_frame(); + if let Some((_, key_name)) = hive_col { + let hive_df = unique_key_frame(hive.df(), &key_name)?; let chunks = get_partitions(&hive_df); @@ -162,35 +174,26 @@ pub fn rewrite_hive( // We do that by pushing down an is_in predicate // Later in the optimizer we prune the hive paths // based on all the predicates. - let mut hive_cols = None; + // Any hive column is valid to split on, prefer the leftmost as it is the coarsest. + let mut hive_cols: Option<(usize, PlSmallStr, PlSmallStr)> = None; let hive_left_schema = hive_left.schema(); let hive_right_schema = hive_right.schema(); for (l, r) in left_on.iter().zip(right_on.iter()) { let l = expr_arena.get(l.node()); let r = expr_arena.get(r.node()); if let (AExpr::Column(l), AExpr::Column(r)) = (l, r) { - if hive_left_schema.index_of(l) == Some(0) - && hive_right_schema.index_of(r) == Some(0) + if let Some(idx) = hive_left_schema.index_of(l) + && hive_right_schema.contains(r) + && hive_cols.as_ref().is_none_or(|(best, _, _)| idx < *best) { - hive_cols = Some((l.clone(), r.clone())); - break; + hive_cols = Some((idx, l.clone(), r.clone())); } } } - if let Some((l, r)) = hive_cols { - let hive_l = hive_left - .df() - .select_at_idx(0) - .unwrap() - .clone() - .into_frame(); - let hive_r = hive_right - .df() - .select_at_idx(0) - .unwrap() - .clone() - .into_frame(); + if let Some((_, l, r)) = hive_cols { + let hive_l = unique_key_frame(hive_left.df(), &l)?; + let hive_r = unique_key_frame(hive_right.df(), &r)?; let partitions = hive_l .join( @@ -199,6 +202,7 @@ pub fn rewrite_hive( [r.as_str()], JoinArgs { how: options.args.how.clone(), + nulls_equal: options.args.nulls_equal, ..Default::default() }, None, @@ -343,7 +347,8 @@ fn make_predicate( AExprBuilder::col(predicate_name, expr_arena) .is_in( AExprBuilder::lit(LiteralValue::Series(SpecialEq::new(values)), expr_arena), - false, // nulls_equal + // Hive __HIVE_DEFAULT_PARTITION__ can produce nulls + true, expr_arena, ) .expr_ir_unnamed() diff --git a/py-polars/tests/unit/io/test_multiscan.py b/py-polars/tests/unit/io/test_multiscan.py index 515b114d761f..1429e02d7b38 100644 --- a/py-polars/tests/unit/io/test_multiscan.py +++ b/py-polars/tests/unit/io/test_multiscan.py @@ -1251,3 +1251,120 @@ def test_hive_join_rewrite_semi_join(tmp_path: Path) -> None: "foo" ), ) + + +@pytest.mark.write_disk +def test_hive_join_rewrite_repeated_partition_values_28617(tmp_path: Path) -> None: + left_root = tmp_path / "left" + right_root = tmp_path / "right" + + # Nested partitioning: the first hive column repeats once per file, so the + # branches must be deduplicated on the partition value. + pl.DataFrame( + { + "cohort_index": [0, 0], + "chunk_index": [1, 2], + "index": [1, 1], + "value": [1, 1], + } + ).write_parquet(left_root, partition_by=["cohort_index", "chunk_index"]) + pl.DataFrame( + { + "cohort_index": [0, 0, 0, 0, 0, 0], + "chunk_index": [1, 1, 1, 2, 2, 2], + "index": [1, 1, 1, 1, 1, 1], + "raster_index": [1, 2, 3, 1, 2, 3], + } + ).write_parquet(right_root, partition_by=["cohort_index", "chunk_index"]) + + left = pl.scan_parquet(left_root, hive_partitioning=True) + right = pl.scan_parquet(right_root, hive_partitioning=True) + + on = ["chunk_index", "cohort_index", "index"] + q = right.join(left, on=on, how="inner") + + out = q.sort("chunk_index", "raster_index").collect() + assert out.height == 6 + assert_frame_equal( + out, + q.collect(optimizations=pl.QueryOptFlags(predicate_pushdown=False)).sort( + "chunk_index", "raster_index" + ), + ) + + +@pytest.mark.write_disk +def test_hive_group_by_rewrite_repeated_partition_values(tmp_path: Path) -> None: + root = tmp_path / "root" + + # `foo` repeats over the nested `bar` partitions, so both `foo` values must + # end up in exactly one branch. + pl.DataFrame( + {"foo": [1, 1, 2, 2], "bar": [1, 2, 1, 2], "x": [1, 2, 3, 4]} + ).write_parquet(root, partition_by=["foo", "bar"]) + + lf = pl.scan_parquet(root, hive_partitioning=True) + q = lf.group_by("foo").agg(pl.sum("x")) + plan = q.explain() + + assert plan.startswith("UNION[maintain_order: false]") + assert plan.count("AGGREGATE") == 2 + + out = q.sort("foo").collect() + assert_frame_equal(out, pl.DataFrame({"foo": [1, 2], "x": [3, 7]})) + + +@pytest.mark.write_disk +def test_hive_group_by_rewrite_null_partition(tmp_path: Path) -> None: + root = tmp_path / "root" + (root / "foo=1").mkdir(parents=True) + (root / "foo=__HIVE_DEFAULT_PARTITION__").mkdir(parents=True) + + pl.DataFrame({"x": [10, 20]}).write_parquet(root / "foo=1" / "data.parquet") + pl.DataFrame({"x": [5]}).write_parquet( + root / "foo=__HIVE_DEFAULT_PARTITION__" / "data.parquet" + ) + + lf = pl.scan_parquet(root, hive_partitioning=True) + q = lf.group_by("foo").agg(pl.sum("x")).sort("foo", nulls_last=True) + + assert_frame_equal( + q.collect(), + pl.DataFrame({"foo": [1, None], "x": [30, 5]}), + ) + assert_frame_equal( + q.collect(), + q.collect(optimizations=pl.QueryOptFlags(predicate_pushdown=False)), + ) + + +@pytest.mark.parametrize("nulls_equal", [False, True]) +@pytest.mark.write_disk +def test_hive_join_rewrite_null_partition(tmp_path: Path, nulls_equal: bool) -> None: + left_root = tmp_path / "left" + right_root = tmp_path / "right" + + for root, name, frames in [ + (left_root, "x", {"foo=1": [10], "foo=2": [20]}), + (right_root, "y", {"foo=1": [100]}), + ]: + for part, values in frames.items(): + (root / part).mkdir(parents=True) + pl.DataFrame({name: values}).write_parquet(root / part / "data.parquet") + + (root / "foo=__HIVE_DEFAULT_PARTITION__").mkdir(parents=True) + pl.DataFrame({name: [0]}).write_parquet( + root / "foo=__HIVE_DEFAULT_PARTITION__" / "data.parquet" + ) + + left = pl.scan_parquet(left_root, hive_partitioning=True) + right = pl.scan_parquet(right_root, hive_partitioning=True) + + for how in ["inner", "left"]: + q = left.join(right, on="foo", how=how, nulls_equal=nulls_equal).sort( # type: ignore[arg-type] + "foo", "x", nulls_last=True + ) + assert_frame_equal( + q.collect(), + q.collect(optimizations=pl.QueryOptFlags(predicate_pushdown=False)), + )