Fix complete groupby

CHANGELOG.md

@@ -1,6 +1,7 @@
 # Changelog
 
 ## [Unreleased]
+-   [PERF] Speed up complete for unsorted domains and many groups. - Issue #1669 @SumanGouda
 -   [PERF] Reduce peak memory in `_build_indexer_reorder_contents` for wide frames (reps >= 8) using single NumPy allocation; tall frames with few repetitions retain the original reshape path. - Issue #1655 @Anupam2400
 -   [ENH] `conditional_join` now picks the most selective `<`/`<=`/`>`/`>=` predicate as its binary-search anchor (instead of the first one supplied) when `keep` is `'first'` or `'last'`, fixing pathological slowdowns from unfavorable predicate ordering; the choice is estimated from a fixed-size sample so the selection cost no longer scales with input size; `keep='all'` output is unaffected. - Issue #1641 @samukweku
 -   [BUG] Fix `conditional_join` crash for `keep='last'` joins with a single `<`/`<=` window and multiple non-equi conditions - a missing `counts` argument to the Rust `index_starts_only_keep_last` call. - Issue #1641 @samukweku

janitor/functions/complete.py

@@ -280,17 +280,24 @@ def _computations_complete(
     A DataFrame, with rows of missing values, if any, is returned.
     """
     check("explicit", explicit, [bool])
-
     check("sort", sort, [bool])
 
     fill_value_check = is_scalar(fill_value), isinstance(fill_value, dict)
     if not any(fill_value_check):
         raise TypeError("fill_value should either be a dictionary or a scalar value.")
+
+    cols = getattr(df, "obj", df).columns
     if fill_value_check[-1]:
-        check_column(df, fill_value)
+        missing_cols = set(fill_value.keys()) - set(cols)
+        if missing_cols:
+            raise ValueError(
+                f"The following columns in `fill_value` are not present in the dataframe: {missing_cols}"
+            )
         for column_name, value in fill_value.items():
             if not is_scalar(value):
                 raise ValueError(f"The value for {column_name} should be a scalar.")
+
+    # Call expand on original df (preserves DataFrameGroupBy behavior)
     uniques = df.expand(*columns, by=by, sort=sort)
     if (by is None) and isinstance(df, pd.DataFrame):
         merge_columns = uniques.columns.tolist()
@@ -304,14 +311,17 @@ def _computations_complete(
         )
         merge_columns = [*uniques.index.names]
         merge_columns.extend(uniques.columns.tolist())
+
     if not isinstance(df, pd.DataFrame):
         df = df.obj
+
     columns = df.columns
     if (fill_value is not None) and not explicit:
         # to get a name that does not exist in the columns
         indicator = "".join(columns)
     else:
         indicator = False
+
     out = pd.merge(
         uniques,
         df,
@@ -320,29 +330,39 @@ def _computations_complete(
         sort=False,
         indicator=indicator,
     )
+
     if indicator:
         indicator = out.pop(indicator)
+
     if not out.columns.equals(columns):
         out = out.reindex(columns=columns)
+
     if fill_value is None:
         return out
+
     # keep only columns that are not part of column_checker
     # IOW, we are excluding columns that were not used
     # to generate the combinations
     null_columns = out.columns.difference(merge_columns)
     null_columns = [col for col in null_columns if out[col].hasnans]
+
     if not null_columns:
         return out
+
     if is_scalar(fill_value):
         # faster when fillna operates on a Series basis
         fill_value = {col: fill_value for col in null_columns}
     else:
-        fill_value = {col: fill_value[col] for col in null_columns}
+        fill_value = {
+            col: fill_value[col] for col in null_columns if col in fill_value
+        }
 
     if not fill_value:
         return out
+
     if explicit:
         return out.fillna(fill_value)
+
     # when explicit is False
     # use the indicator parameter to identify rows
     # for `left_only`, and fill the relevant columns

tests/functions/test_complete.py

@@ -809,3 +809,83 @@ def test_MI_1(MI):
     ).rename_axis(columns=[None, None])
     actual = MI.iloc[:2].complete({("a", "bar"): pd.Series(range(1, 5))})
     assert_frame_equal(actual, expected)
+
+
+@pytest.fixture
+def sample_df():
+    return pd.DataFrame(
+        {
+            "state": ["CA", "CA", "NY"],
+            "year": [2020, 2021, 2020],
+            "sales": [100, 150, 200],
+            "profit": [10, 15, 20],
+        }
+    )
+
+
+def test_complete_invalid_fill_value_column(sample_df):
+    with pytest.raises(
+        ValueError, match="not present in the dataframe"
+    ):
+        sample_df.complete(
+            "state", "year", fill_value={"non_existent_col": 0}
+        )
+
+
+def test_complete_partial_dict_fill(sample_df):
+    # Only specifying fill value for 'sales', leaving 'profit' alone
+    result = sample_df.complete(
+        "state", "year", fill_value={"sales": 0}
+    )
+
+    expected = pd.DataFrame(
+        {
+            "state": ["CA", "CA", "NY", "NY"],
+            "year": [2020, 2021, 2020, 2021],
+            "sales": [100.0, 150.0, 200.0, 0.0],
+            "profit": [10.0, 15.0, 20.0, None],
+        }
+    )
+    assert_frame_equal(result, expected)
+
+
+def test_complete_groupby(sample_df):
+    # To expand NY to include 2021, explicit categories/sequence must be passed,
+    # or group-by completes existing combinations within each group.
+    result = sample_df.groupby("state").complete(
+        {"year": [2020, 2021]}, fill_value=0
+    )
+
+    assert isinstance(result, pd.DataFrame)
+    # Now both CA and NY get 2020 & 2021 -> Total 4 rows
+    assert len(result) == 4
+
+    # Verify NY 2021 missing row was filled with 0
+    ny_2021 = result[(result["state"] == "NY") & (result["year"] == 2021)].iloc[0]
+    assert ny_2021["sales"] == 0
+    assert ny_2021["profit"] == 0
+
+
+def test_complete_explicit_flag(sample_df):
+    df_with_nan = sample_df.copy()
+    df_with_nan.loc[0, "sales"] = None  # Existing NaN in original data
+
+    # explicit=False should only fill implicit missing combinations, not existing NaNs
+    res_implicit = df_with_nan.complete(
+        "state", "year", fill_value={"sales": 0}, explicit=False
+    )
+    assert pd.isna(res_implicit.loc[0, "sales"])  # Original NaN stays NaN
+    assert res_implicit.loc[3, "sales"] == 0  # Missing row gets filled
+
+    # explicit=True fills both original NaNs and missing combinations
+    res_explicit = df_with_nan.complete(
+        "state", "year", fill_value={"sales": 0}, explicit=True
+    )
+    assert res_explicit.loc[0, "sales"] == 0
+    assert res_explicit.loc[3, "sales"] == 0
+
+
+def test_complete_scalar_fill(sample_df):
+    result = sample_df.complete("state", "year", fill_value=0)
+    assert result.loc[3, "sales"] == 0
+    assert result.loc[3, "profit"] == 0