fix: Fix duplicate hive values in hive pre-partition rewrite

crates/polars-plan/src/plans/optimizer/predicate_pushdown/hive.rs

@@ -44,6 +44,16 @@ fn get_partitions(hive_df: &DataFrame) -> Vec<DataFrame> {
     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<DataFrame> {
+    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()

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)),
+        )