Repository: pyjanitor-devs/pyjanitor
Pull request #1389 (2024-07-13 14:52:22)

[ENH] improve performance for polars' `complete`
# PR Description

Please describe the changes proposed in the pull request:

- remove ambiguity for multiple columns - user creates a struct before calling complete
- update code to remove performance warning ( uses `collect_schema` instead of `columns`)
-

<!-- Doing so provides maintainers with context on what the PR is, and can help us more effectively review your PR. -->

**This PR relates to #1352 .**

<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [x] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.md`.
<!-- We'd like to acknowledge your contributions! -->
3. [ ] Add a line to `CHANGELOG.md` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

# Automatic checks

There will be automatic checks run on the PR. These include:

- Building a preview of the docs on Netlify
- Automatically linting the code
- Making sure the code is documented
- Making sure that all tests are passed
- Making sure that code coverage doesn't go down.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


ericmjl: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY -->
🚀 Deployed on https://deploy-preview-1389--pyjanitor.netlify.app
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyjanitor-devs/pyjanitor/pull/1389?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs) Report
Attention: Patch coverage is `96.96970%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 88.39%. Comparing base [(`62c57c6`)](https://app.codecov.io/gh/pyjanitor-devs/pyjanitor/commit/62c57c6ad27196a10e017f710ce7d523bc2541ba?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs) to head [(`911abd8`)](https://app.codecov.io/gh/pyjanitor-devs/pyjanitor/commit/911abd8faf51f861dfedbb041b5071806e55ec29?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs).
> Report is 33 commits behind head on dev.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##              dev    #1389      +/-   ##
==========================================
- Coverage   94.48%   88.39%   -6.09%     
==========================================
  Files          80       87       +7     
  Lines        4367     5352     +985     
==========================================
+ Hits         4126     4731     +605     
- Misses        241      621     +380     
```

</details>
ericmjl: Looking great, @samukweku! Please do the honors of hitting the merge button!

Patch:
diff --git a/environment-dev.yml b/environment-dev.yml
index 2543e2c..373ce61 100644
--- a/environment-dev.yml
+++ b/environment-dev.yml
@@ -25,7 +25,7 @@ dependencies:
   - mypy
   - natsort
   - numba
-  - numpy
+  - numpy==1.24.4
   - openpyxl
   - pandas-flavor
   - pandas-vet
@@ -51,3 +51,5 @@ dependencies:
   - xorg-libxrender
   - pip:
       - mknotebooks
+  # Temporarily pinned to fix CI
+  - setuptools==70.3.0
diff --git a/janitor/functions/convert_date.py b/janitor/functions/convert_date.py
index fbd1974..eac2053 100644
--- a/janitor/functions/convert_date.py
+++ b/janitor/functions/convert_date.py
@@ -102,7 +102,6 @@ def convert_matlab_date(
     return df.assign(**dictionary)
 
 
-@pf.register_dataframe_method
 @pf.register_dataframe_method
 @refactored_function(
     message=(
diff --git a/janitor/polars/__init__.py b/janitor/polars/__init__.py
index 972db2c..ff0fd1b 100644
--- a/janitor/polars/__init__.py
+++ b/janitor/polars/__init__.py
@@ -1,5 +1,5 @@
 from .clean_names import clean_names, make_clean_names
-from .complete import complete
+from .complete import complete, expand
 from .dates_to_polars import convert_excel_date, convert_matlab_date
 from .pivot_longer import pivot_longer, pivot_longer_spec
 from .row_to_names import row_to_names
@@ -10,6 +10,7 @@ __all__ = [
     "clean_names",
     "make_clean_names",
     "row_to_names",
+    "expand",
     "complete",
     "convert_excel_date",
     "convert_matlab_date",
diff --git a/janitor/polars/complete.py b/janitor/polars/complete.py
index ddd6f0a..489b19a 100644
--- a/janitor/polars/complete.py
+++ b/janitor/polars/complete.py
@@ -21,6 +21,250 @@ except ImportError:
     )
 
 
+@register_lazyframe_method
+@register_dataframe_method
+def expand(
+    df: pl.DataFrame | pl.LazyFrame,
+    *columns: ColumnNameOrSelector,
+    sort: bool = False,
+    by: ColumnNameOrSelector = None,
+) -> pl.DataFrame | pl.LazyFrame:
+    """
+    Creates a DataFrame from a cartesian combination of all inputs.
+
+    Inspiration is from tidyr's expand() function.
+
+    expand() is often useful with
+    [pl.DataFrame.join](https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.join.html)
+    to convert implicit
+    missing values to explicit missing values - similar to
+    [`complete`][janitor.polars.complete.complete].
+
+    It can also be used to figure out which combinations are missing
+    (e.g identify gaps in your DataFrame).
+
+    The variable `columns` parameter can be a string,
+    a ColumnSelector, a polars expression, or a polars Series.
+
+    `expand` can also be applied to a LazyFrame.
+
+    Examples:
+        >>> import polars as pl
+        >>> import janitor.polars
+        >>> data = [{'type': 'apple', 'year': 2010, 'size': 'XS'},
+        ...         {'type': 'orange', 'year': 2010, 'size': 'S'},
+        ...         {'type': 'apple', 'year': 2012, 'size': 'M'},
+        ...         {'type': 'orange', 'year': 2010, 'size': 'S'},
+        ...         {'type': 'orange', 'year': 2011, 'size': 'S'},
+        ...         {'type': 'orange', 'year': 2012, 'size': 'M'}]
+        >>> df = pl.DataFrame(data)
+        >>> df
+        shape: (6, 3)
+        ┌────────┬──────┬──────┐
+        │ type   ┆ year ┆ size │
+        │ ---    ┆ ---  ┆ ---  │
+        │ str    ┆ i64  ┆ str  │
+        ╞════════╪══════╪══════╡
+        │ apple  ┆ 2010 ┆ XS   │
+        │ orange ┆ 2010 ┆ S    │
+        │ apple  ┆ 2012 ┆ M    │
+        │ orange ┆ 2010 ┆ S    │
+        │ orange ┆ 2011 ┆ S    │
+        │ orange ┆ 2012 ┆ M    │
+        └────────┴──────┴──────┘
+
+        Get unique observations:
+        >>> df.expand('type',sort=True)
+        shape: (2, 1)
+        ┌────────┐
+        │ type   │
+        │ ---    │
+        │ str    │
+        ╞════════╡
+        │ apple  │
+        │ orange │
+        └────────┘
+        >>> df.expand('size',sort=True)
+        shape: (3, 1)
+        ┌──────┐
+        │ size │
+        │ ---  │
+        │ str  │
+        ╞══════╡
+        │ M    │
+        │ S    │
+        │ XS   │
+        └──────┘
+        >>> df.expand('type', 'size',sort=True)
+        shape: (6, 2)
+        ┌────────┬──────┐
+        │ type   ┆ size │
+        │ ---    ┆ ---  │
+        │ str    ┆ str  │
+        ╞════════╪══════╡
+        │ apple  ┆ M    │
+        │ apple  ┆ S    │
+        │ apple  ┆ XS   │
+        │ orange ┆ M    │
+        │ orange ┆ S    │
+        │ orange ┆ XS   │
+        └────────┴──────┘
+        >>> with pl.Config(tbl_rows=-1):
+        ...     df.expand('type','size','year',sort=True)
+        shape: (18, 3)
+        ┌────────┬──────┬──────┐
+        │ type   ┆ size ┆ year │
+        │ ---    ┆ ---  ┆ ---  │
+        │ str    ┆ str  ┆ i64  │
+        ╞════════╪══════╪══════╡
+        │ apple  ┆ M    ┆ 2010 │
+        │ apple  ┆ M    ┆ 2011 │
+        │ apple  ┆ M    ┆ 2012 │
+        │ apple  ┆ S    ┆ 2010 │
+        │ apple  ┆ S    ┆ 2011 │
+        │ apple  ┆ S    ┆ 2012 │
+        │ apple  ┆ XS   ┆ 2010 │
+        │ apple  ┆ XS   ┆ 2011 │
+        │ apple  ┆ XS   ┆ 2012 │
+        │ orange ┆ M    ┆ 2010 │
+        │ orange ┆ M    ┆ 2011 │
+        │ orange ┆ M    ┆ 2012 │
+        │ orange ┆ S    ┆ 2010 │
+        │ orange ┆ S    ┆ 2011 │
+        │ orange ┆ S    ┆ 2012 │
+        │ orange ┆ XS   ┆ 2010 │
+        │ orange ┆ XS   ┆ 2011 │
+        │ orange ┆ XS   ┆ 2012 │
+        └────────┴──────┴──────┘
+
+        Get observations that only occur in the data:
+        >>> df.expand(pl.struct('type','size'),sort=True).unnest('type')
+        shape: (4, 2)
+        ┌────────┬──────┐
+        │ type   ┆ size │
+        │ ---    ┆ ---  │
+        │ str    ┆ str  │
+        ╞════════╪══════╡
+        │ apple  ┆ M    │
+        │ apple  ┆ XS   │
+        │ orange ┆ M    │
+        │ orange ┆ S    │
+        └────────┴──────┘
+        >>> df.expand(pl.struct('type','size','year'),sort=True).unnest('type')
+        shape: (5, 3)
+        ┌────────┬──────┬──────┐
+        │ type   ┆ size ┆ year │
+        │ ---    ┆ ---  ┆ ---  │
+        │ str    ┆ str  ┆ i64  │
+        ╞════════╪══════╪══════╡
+        │ apple  ┆ M    ┆ 2012 │
+        │ apple  ┆ XS   ┆ 2010 │
+        │ orange ┆ M    ┆ 2012 │
+        │ orange ┆ S    ┆ 2010 │
+        │ orange ┆ S    ┆ 2011 │
+        └────────┴──────┴──────┘
+
+        Expand the DataFrame to include new observations:
+        >>> with pl.Config(tbl_rows=-1):
+        ...     df.expand('type','size',pl.int_range(2010,2014).alias('new_year'),sort=True)
+        shape: (24, 3)
+        ┌────────┬──────┬──────────┐
+        │ type   ┆ size ┆ new_year │
+        │ ---    ┆ ---  ┆ ---      │
+        │ str    ┆ str  ┆ i64      │
+        ╞════════╪══════╪══════════╡
+        │ apple  ┆ M    ┆ 2010     │
+        │ apple  ┆ M    ┆ 2011     │
+        │ apple  ┆ M    ┆ 2012     │
+        │ apple  ┆ M    ┆ 2013     │
+        │ apple  ┆ S    ┆ 2010     │
+        │ apple  ┆ S    ┆ 2011     │
+        │ apple  ┆ S    ┆ 2012     │
+        │ apple  ┆ S    ┆ 2013     │
+        │ apple  ┆ XS   ┆ 2010     │
+        │ apple  ┆ XS   ┆ 2011     │
+        │ apple  ┆ XS   ┆ 2012     │
+        │ apple  ┆ XS   ┆ 2013     │
+        │ orange ┆ M    ┆ 2010     │
+        │ orange ┆ M    ┆ 2011     │
+        │ orange ┆ M    ┆ 2012     │
+        │ orange ┆ M    ┆ 2013     │
+        │ orange ┆ S    ┆ 2010     │
+        │ orange ┆ S    ┆ 2011     │
+        │ orange ┆ S    ┆ 2012     │
+        │ orange ┆ S    ┆ 2013     │
+        │ orange ┆ XS   ┆ 2010     │
+        │ orange ┆ XS   ┆ 2011     │
+        │ orange ┆ XS   ┆ 2012     │
+        │ orange ┆ XS   ┆ 2013     │
+        └────────┴──────┴──────────┘
+
+        Filter for missing observations:
+        >>> columns = ('type','size','year')
+        >>> with pl.Config(tbl_rows=-1):
+        ...     df.expand(*columns).join(df, how='anti', on=columns).sort(by=pl.all())
+        shape: (13, 3)
+        ┌────────┬──────┬──────┐
+        │ type   ┆ size ┆ year │
+        │ ---    ┆ ---  ┆ ---  │
+        │ str    ┆ str  ┆ i64  │
+        ╞════════╪══════╪══════╡
+        │ apple  ┆ M    ┆ 2010 │
+        │ apple  ┆ M    ┆ 2011 │
+        │ apple  ┆ S    ┆ 2010 │
+        │ apple  ┆ S    ┆ 2011 │
+        │ apple  ┆ S    ┆ 2012 │
+        │ apple  ┆ XS   ┆ 2011 │
+        │ apple  ┆ XS   ┆ 2012 │
+        │ orange ┆ M    ┆ 2010 │
+        │ orange ┆ M    ┆ 2011 │
+        │ orange ┆ S    ┆ 2012 │
+        │ orange ┆ XS   ┆ 2010 │
+        │ orange ┆ XS   ┆ 2011 │
+        │ orange ┆ XS   ┆ 2012 │
+        └────────┴──────┴──────┘
+
+        Expand within each group, using `by`:
+        >>> with pl.Config(tbl_rows=-1):
+        ...     df.expand('year','size',by='type',sort=True)
+        shape: (10, 3)
+        ┌────────┬──────┬──────┐
+        │ type   ┆ year ┆ size │
+        │ ---    ┆ ---  ┆ ---  │
+        │ str    ┆ i64  ┆ str  │
+        ╞════════╪══════╪══════╡
+        │ apple  ┆ 2010 ┆ M    │
+        │ apple  ┆ 2010 ┆ XS   │
+        │ apple  ┆ 2012 ┆ M    │
+        │ apple  ┆ 2012 ┆ XS   │
+        │ orange ┆ 2010 ┆ M    │
+        │ orange ┆ 2010 ┆ S    │
+        │ orange ┆ 2011 ┆ M    │
+        │ orange ┆ 2011 ┆ S    │
+        │ orange ┆ 2012 ┆ M    │
+        │ orange ┆ 2012 ┆ S    │
+        └────────┴──────┴──────┘
+
+    !!! info "New in version 0.28.0"
+
+    Args:
+        *columns: This refers to the columns to be completed.
+            It can be a string or a column selector or a polars expression.
+            A polars expression can be used to introduced new values,
+            as long as the polars expression has a name that already exists
+            in the DataFrame.
+        sort: Sort the DataFrame based on *columns.
+        by: Column(s) to group by.
+
+    Returns:
+        A polars DataFrame/LazyFrame.
+    """  # noqa: E501
+    if not columns:
+        return df
+    uniques, _ = _expand(df=df, columns=columns, by=by, sort=sort)
+    return uniques
+
+
 @register_lazyframe_method
 @register_dataframe_method
 def complete(
@@ -38,19 +282,11 @@ def complete(
     In a way, it is the inverse of `pl.drop_nulls`,
     as it exposes implicitly missing rows.
 
-    If the combination involves multiple columns, pass it as a struct,
-    with an alias - the name of the struct should not exist in the DataFrame.
-
     If new values need to be introduced, a polars Expression
-    with the new values can be passed, as long as the polars Expression
+    or a polars Series with the new values can be passed,
+    as long as the polars Expression/Series
     has a name that already exists in the DataFrame.
 
-    It is up to the user to ensure that the polars expression returns
-    unique values and/or sorted values.
-
-    Note that if the polars expression evaluates to a struct,
-    then the fields, not the name, should already exist in the DataFrame.
-
     `complete` can also be applied to a LazyFrame.
 
     Examples:
@@ -105,14 +341,10 @@ def complete(
 
         Cross all possible `group` values with the unique pairs of
         `(item_id, item_name)` that already exist in the data.
-        For such situations, where there is a group of columns,
-        pass it in as a struct:
         >>> with pl.Config(tbl_rows=-1):
-        ...     df.complete(
-        ...         "group",
-        ...         pl.struct("item_id", "item_name").unique().sort().alias("rar"),
-        ...         sort=True
-        ...     )
+        ...     df.select(
+        ...         "group", pl.struct("item_id", "item_name"), "value1", "value2"
+        ...     ).complete("group", "item_id", sort=True).unnest("item_id")
         shape: (8, 5)
         ┌───────┬─────────┬───────────┬────────┬────────┐
         │ group ┆ item_id ┆ item_name ┆ value1 ┆ value2 │
@@ -131,13 +363,15 @@ def complete(
 
         Fill in nulls:
         >>> with pl.Config(tbl_rows=-1):
-        ...     df.complete(
+        ...     df.select(
+        ...         "group", pl.struct("item_id", "item_name"), "value1", "value2"
+        ...     ).complete(
         ...         "group",
-        ...         pl.struct("item_id", "item_name").unique().sort().alias('rar'),
+        ...         "item_id",
         ...         fill_value={"value1": 0, "value2": 99},
         ...         explicit=True,
         ...         sort=True,
-        ...     )
+        ...     ).unnest("item_id")
         shape: (8, 5)
         ┌───────┬─────────┬───────────┬────────┬────────┐
         │ group ┆ item_id ┆ item_name ┆ value1 ┆ value2 │
@@ -155,15 +389,17 @@ def complete(
         └───────┴─────────┴───────────┴────────┴────────┘
 
         Limit the fill to only the newly created
-        missing values with `explicit = FALSE`
+        missing values with `explicit = FALSE`:
         >>> with pl.Config(tbl_rows=-1):
-        ...     df.complete(
+        ...     df.select(
+        ...         "group", pl.struct("item_id", "item_name"), "value1", "value2"
+        ...     ).complete(
         ...         "group",
-        ...         pl.struct("item_id", "item_name").unique().sort().alias('rar'),
+        ...         "item_id",
         ...         fill_value={"value1": 0, "value2": 99},
         ...         explicit=False,
         ...         sort=True,
-        ...     )
+        ...     ).unnest("item_id").sort(pl.all())
         shape: (8, 5)
         ┌───────┬─────────┬───────────┬────────┬────────┐
         │ group ┆ item_id ┆ item_name ┆ value1 ┆ value2 │
@@ -295,8 +531,6 @@ def complete(
             A polars expression can be used to introduced new values,
             as long as the polars expression has a name that already exists
             in the DataFrame.
-            It is up to the user to ensure that the polars expression returns
-            unique values.
         fill_value: Scalar value or polars expression to use instead of nulls
             for missing combinations. A dictionary, mapping columns names
             to a scalar value is also accepted.
@@ -311,6 +545,8 @@ def complete(
     Returns:
         A polars DataFrame/LazyFrame.
     """  # noqa: E501
+    if not columns:
+        return df
     return _complete(
         df=df,
         columns=columns,
@@ -321,11 +557,9 @@ def complete(
     )
 
 
-def _complete(
+def _expand(
     df: pl.DataFrame | pl.LazyFrame,
     columns: tuple[ColumnNameOrSelector],
-    fill_value: dict | Any | pl.Expr,
-    explicit: bool,
     sort: bool,
     by: ColumnNameOrSelector,
 ) -> pl.DataFrame | pl.LazyFrame:
@@ -334,78 +568,87 @@ def _complete(
 
     A DataFrame, with rows of missing values, if any, is returned.
     """
-    if not columns:
-        return df
 
     check("sort", sort, [bool])
-    check("explicit", explicit, [bool])
     _columns = []
     for column in columns:
         if isinstance(column, str):
             col = pl.col(column).unique()
-            if sort:
-                col = col.sort()
-            _columns.append(col)
         elif cs.is_selector(column):
             col = column.as_expr().unique()
-            if sort:
-                col = col.sort()
-            _columns.append(col)
         elif isinstance(column, pl.Expr):
-            _columns.append(column)
+            col = column.unique()
+        elif isinstance(column, pl.Series):
+            col = column.unique()
         else:
             raise TypeError(
                 f"The argument passed to the columns parameter "
                 "should either be a string, a column selector, "
-                "or a polars expression, instead got - "
-                f"{type(column)}."
+                "a polars expression, or a polars Series; "
+                f"instead got - {type(column)}."
             )
+        if sort:
+            col = col.sort()
+        _columns.append(col)
     by_does_not_exist = by is None
     if by_does_not_exist:
         _columns = [column.implode() for column in _columns]
         uniques = df.select(_columns)
-        _columns = uniques.columns
+        uniques_schema = uniques.collect_schema()
+        _columns = uniques_schema.names()
     else:
         uniques = df.group_by(by, maintain_order=sort).agg(_columns)
-        _by = uniques.select(by).columns
-        _columns = uniques.select(pl.exclude(_by)).columns
+        uniques_schema = uniques.collect_schema()
+        _columns = cs.expand_selector(
+            uniques_schema, cs.exclude(by), strict=False
+        )
     for column in _columns:
         uniques = uniques.explode(column)
 
-    _columns = [
-        column
-        for column, dtype in zip(_columns, uniques.select(_columns).dtypes)
-        # this way we ensure there is no tampering with existing struct columns
-        if (dtype == pl.Struct) and (column not in df.columns)
-    ]
+    return uniques, uniques_schema
 
-    if _columns:
-        for column in _columns:
-            uniques = uniques.unnest(columns=column)
 
-    no_columns_to_fill = set(df.columns) == set(uniques.columns)
-    if fill_value is None or no_columns_to_fill:
-        return uniques.join(df, on=uniques.columns, how="left", coalesce=True)
+def _complete(
+    df: pl.DataFrame | pl.LazyFrame,
+    columns: tuple[ColumnNameOrSelector],
+    fill_value: dict | Any | pl.Expr,
+    explicit: bool,
+    sort: bool,
+    by: ColumnNameOrSelector,
+) -> pl.DataFrame | pl.LazyFrame:
+    """
+    This function computes the final output for the `complete` function.
+
+    A DataFrame, with rows of missing values, if any, is returned.
+    """
+    check("explicit", explicit, [bool])
+    uniques, uniques_schema = _expand(df=df, columns=columns, sort=sort, by=by)
+
+    df_columns = df.collect_schema()
+    columns_to_fill = df_columns.keys() ^ uniques_schema.keys()
+    if (fill_value is None) or not columns_to_fill:
+        return uniques.join(
+            df, on=uniques_schema.names(), how="left", coalesce=True
+        )
     idx = None
-    columns_to_select = df.columns
+    columns_to_select = df_columns.names()
     if not explicit:
-        idx = "".join(df.columns)
+        idx = "".join(columns_to_select)
         idx = f"{idx}_"
         df = df.with_row_index(name=idx)
-    df = uniques.join(df, on=uniques.columns, how="left", coalesce=True)
+    df = uniques.join(df, on=uniques_schema.names(), how="left", coalesce=True)
     # exclude columns that were not used
     # to generate the combinations
-    exclude_columns = uniques.columns
+    exclude_columns = uniques_schema.names()
     if idx:
         exclude_columns.append(idx)
-    _columns = [
-        column for column in columns_to_select if column not in exclude_columns
-    ]
+    _columns = set(columns_to_select).difference(exclude_columns)
+
     if isinstance(fill_value, dict):
+        _columns = _columns.intersection(fill_value)
         fill_value = [
-            pl.col(column_name).fill_null(value=value)
-            for column_name, value in fill_value.items()
-            if column_name in _columns
+            pl.col(column_name).fill_null(value=fill_value[column_name])
+            for column_name in _columns
         ]
     else:
         fill_value = [


Test patch:
diff --git a/tests/polars/functions/test_complete_polars.py b/tests/polars/functions/test_complete_polars.py
index 4af7e20..5fd36c6 100644
--- a/tests/polars/functions/test_complete_polars.py
+++ b/tests/polars/functions/test_complete_polars.py
@@ -47,7 +47,7 @@ def test_empty_groups(fill_df):
     """Raise TypeError if wrong column type is passed."""
     msg = "The argument passed to the columns parameter "
     msg += "should either be a string, a column selector, "
-    msg += "or a polars expression, instead got.+"
+    msg += "a polars expression, or a polars Series; instead got.+"
     with pytest.raises(TypeError, match=msg):
         fill_df.complete("group", {})
 
@@ -69,12 +69,16 @@ def test_complete_1(fill_df):
     Test output for janitor.complete.
     """
     trimmed = fill_df.lazy().select(~cs.starts_with("value"))
-    result = trimmed.complete(
-        cs.by_name("group"),
-        pl.struct("item_id", "item_name").alias("rar").unique().sort(),
-        fill_value=0,
-        explicit=False,
-        sort=True,
+    result = (
+        trimmed.select("group", pl.struct("item_id", "item_name"))
+        .complete(
+            cs.by_name("group"),
+            "item_id",
+            fill_value=0,
+            explicit=False,
+            sort=True,
+        )
+        .unnest("item_id")
     )
     expected = pl.DataFrame(
         [
@@ -147,12 +151,18 @@ def test_groupby_complete():
 # https://tidyr.tidyverse.org/reference/complete.html
 def test_complete_2(fill_df):
     """Test output for janitor.complete."""
-    result = fill_df.complete(
-        "group",
-        pl.struct("item_id", "item_name").alias("rar").unique().sort(),
-        fill_value={"value1": 0, "value2": 99},
-        explicit=False,
-        sort=True,
+    result = (
+        fill_df.select(
+            "group", pl.struct("item_id", "item_name"), "value1", "value2"
+        )
+        .complete(
+            "group",
+            pl.col("item_id").unique().sort(),
+            fill_value={"value1": 0, "value2": 99},
+            explicit=False,
+            sort=True,
+        )
+        .unnest("item_id")
     )
     expected = pl.DataFrame(
         [
@@ -241,12 +251,21 @@ def test_complete_multiple_groupings():
         }
     )
 
-    result = df3.complete(
-        pl.struct("meta", "domain1").alias("bar").unique().sort(),
-        pl.struct("project_id", "question_count").alias("foo").unique().sort(),
-        fill_value={"tag_count": 0},
-        sort=True,
-    ).select("project_id", "meta", "domain1", "question_count", "tag_count")
+    result = (
+        df3.select(
+            pl.struct("meta", "domain1"),
+            pl.struct("project_id", "question_count"),
+            "tag_count",
+        )
+        .complete(
+            pl.col("meta").unique().sort(),
+            "project_id",
+            fill_value={"tag_count": 0},
+            sort=True,
+        )
+        .unnest("meta", "project_id")
+        .select("project_id", "meta", "domain1", "question_count", "tag_count")
+    )
     assert_frame_equal(result, output3)
 
 
diff --git a/tests/polars/functions/test_expand_polars.py b/tests/polars/functions/test_expand_polars.py
new file mode 100644
index 0000000..0c69487
--- /dev/null
+++ b/tests/polars/functions/test_expand_polars.py
@@ -0,0 +1,69 @@
+import polars as pl
+import pytest
+from polars.testing import assert_frame_equal
+
+import janitor.polars  # noqa: F401
+
+
+@pytest.fixture
+def df():
+    """pytest fixture"""
+    return pl.DataFrame(
+        dict(
+            group=(1, 2, 1, 2),
+            item_id=(1, 2, 2, 3),
+            item_name=("a", "a", "b", "b"),
+            value1=(1, None, 3, 4),
+            value2=range(4, 8),
+        )
+    )
+
+
+def test_column_None(df):
+    """Test output if *columns is empty."""
+    assert_frame_equal(df.expand(), df)
+
+
+def test_empty_groups(df):
+    """Raise TypeError if wrong column type is passed."""
+    msg = "The argument passed to the columns parameter "
+    msg += "should either be a string, a column selector, "
+    msg += "a polars expression, or a polars Series; instead got.+"
+    with pytest.raises(TypeError, match=msg):
+        df.complete("group", {})
+
+
+def test_type_sort(df):
+    """Raise TypeError if `sort` is not boolean."""
+    with pytest.raises(TypeError):
+        df.complete("group", "item_id", sort=11)
+
+
+def test_expand_1(df):
+    """
+    Test output for janitor.expand.
+    """
+    expected = df.expand("group", "item_id", "item_name", sort=True)
+    actual = (
+        df.select(pl.col("group").unique())
+        .join(df.select(pl.col("item_id").unique()), how="cross")
+        .join(df.select(pl.col("item_name").unique()), how="cross")
+        .sort(by=pl.all())
+    )
+    assert_frame_equal(actual, expected)
+
+
+def test_expand_2(df):
+    """
+    Test output for janitor.expand.
+    """
+    expected = df.expand(
+        "group", df.get_column("item_id"), "item_name", sort=True
+    )
+    actual = (
+        df.select(pl.col("group").unique())
+        .join(df.select(pl.col("item_id").unique()), how="cross")
+        .join(df.select(pl.col("item_name").unique()), how="cross")
+        .sort(by=pl.all())
+    )
+    assert_frame_equal(actual, expected)