Polars unique() behavior

[rjavadi]

@AdrianSosic You alluded in the meeting that Polar's unique() does what we need.
According to the example in the pydocs:

"""Constraint class for excluding entries where occurring labels are not unique.

    This can be useful to remove entries that arise from e.g. a permutation invariance
    as for instance here:

    - A,B,C,D would remain
    - A,A,B,C would be removed
    - A,A,B,B would be removed
    - A,A,B,A would be removed
    - A,C,A,C would be removed
    - A,C,B,C would be removed
    """

But the output of unique() is different:

data = {
    "a": ["A", "A", "A", "A", "A", "A"],
    "b": ["B", "C", "A", "A", "D", "A"],
    "c": ["C", "C", "B", "A", "C", "B"],
    "d": ["D", "D", "C", "B", "B", "A"]
}

df = pl.DataFrame(data)
df = df.unique(maintain_order=True)

Output:

┌─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ str │
╞═════╪═════╪═════╪═════╡
│ A   ┆ B   ┆ C   ┆ D   │
│ A   ┆ C   ┆ C   ┆ D   │
│ A   ┆ A   ┆ B   ┆ C   │
│ A   ┆ A   ┆ A   ┆ B   │
│ A   ┆ D   ┆ C   ┆ B   │
│ A   ┆ A   ┆ B   ┆ A   │
└─────┴─────┴─────┴─────┘

Another example with subset columns:

df = df.unique(subset=["a"], maintain_order=True)

┌─────┬─────┬─────┬─────┐
│ a   ┆ b   ┆ c   ┆ d   │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ str ┆ str │
╞═════╪═════╪═════╪═════╡
│ A   ┆ B   ┆ C   ┆ D   │
└─────┴─────┴─────┴─────┘

It only filters out the rows that are duplicate of each other in the dataframe and I believe it's not what we intend to do. Please let me if I'm wrong.

[AdrianSosic]

HI @rjavadi, sorry, I did not get the context right during our call. unique indeed looks for duplicates along columns, while in the case you want to handle (<-- I missed that context) you need to search for duplicates along rows. So the method indeed won't help you here. Can you again post what was your solution? We should only try to avoid using UDFs and rely only on polars internals – if possible – for performance reasons

[rjavadi]

This is my query:

df = df.with_columns(
    pl.concat_list(pl.col("*")).list.eval(pl.element().n_unique()).explode().alias("temp")
)
# if the value of temp is not equal with number of columns, it means that there are duplicates
filtered_df = df.filter(pl.col('temp') == df.width-1)
print(df)
print(filtered_df)

Output:

df with count of unique values in each row:
shape: (6, 5)
┌─────┬─────┬─────┬─────┬──────┐
│ a   ┆ b   ┆ c   ┆ d   ┆ temp │
│ --- ┆ --- ┆ --- ┆ --- ┆ ---  │
│ str ┆ str ┆ str ┆ str ┆ u32  │
╞═════╪═════╪═════╪═════╪══════╡
│ A   ┆ B   ┆ C   ┆ D   ┆ 4    │
│ A   ┆ C   ┆ C   ┆ D   ┆ 3    │
│ A   ┆ A   ┆ B   ┆ C   ┆ 3    │
│ A   ┆ A   ┆ A   ┆ B   ┆ 2    │
│ A   ┆ D   ┆ C   ┆ B   ┆ 4    │
│ A   ┆ A   ┆ B   ┆ A   ┆ 2    │
└─────┴─────┴─────┴─────┴──────┘

filtered df:

shape: (2, 5)
┌─────┬─────┬─────┬─────┬──────┐
│ a   ┆ b   ┆ c   ┆ d   ┆ temp │
│ --- ┆ --- ┆ --- ┆ --- ┆ ---  │
│ str ┆ str ┆ str ┆ str ┆ u32  │
╞═════╪═════╪═════╪═════╪══════╡
│ A   ┆ B   ┆ C   ┆ D   ┆ 4    │
│ A   ┆ D   ┆ C   ┆ B   ┆ 4    │
└─────┴─────┴─────┴─────┴──────┘

And here's Polar's explain() when using lazy frames:

FILTER [(col("temp")) == (4)] FROM
   WITH_COLUMNS:
   [col("a").list.concat([col("b"), col("c"), col("d")]).eval().explode().alias("temp")], [] 
    DF ["a", "b", "c", "d"]; PROJECT */4 COLUMNS; SELECTION: None

[rjavadi]

Accoring to Polars docs list operations use query engine and have a good performance:

We can apply any Polars operations on the elements of the list with the list.eval (list().eval in Rust) expression! These expressions run entirely on Polars' query engine and can run in parallel, so will be well optimized.