Polars equivalent to SQL `COUNT(DISTINCT expr,[expr...])`, or other method of checking uniqueness - python-polars

Question

When processing data, I often add a check after each step to validate that the data still has the unique key I think it does. For example, I might check that my data is still unique on (a, b). To accomplish this, I would typically check that the number of distinct combinations of columns a and b equals the total number of rows.
In polars, to get a COUNT(DISTINCT ...) I can do
(
df
.select(['a', 'b'])
.unique()
.height
)
But height does not work on LazyFrames, so I need to actually materialize the entire data with this method, I think (?). Is there a better way?
For reference, in R's data.table library I would do
mtc_dt <- data.table::as.data.table(mtcars)
stopifnot(data.table::uniqueN(mtc_dt[, .(mpg, disp)]) == nrow(mtc_dt))
To any contributors reading:
Thanks for the great package! Has sped up many of my workflows to a fraction of the time.

Answer 1

You can use a map function that asserts on the unique count.
This allows you to get an eager DataFrame in the middle of a query plan.
Note that we turn off projection_pushdown optimization, as the optimizer is not able to know which subset of columns we select.
df = pl.DataFrame({
"foo": [1, 2, 3],
"bar": [None, "hello", None]
})
def unique_check(df: pl.DataFrame, subset: list[str]) -> pl.DataFrame:
assert df.select(pl.struct(subset).unique().count()).item() == df.height
return df
out = (df.lazy()
.map(lambda df: unique_check(df, ["foo", "bar"]), projection_pushdown=False)
.select("bar")
.collect()
)
print(out)
shape: (3, 1)
┌───────┐
│ bar │
│ --- │
│ str │
╞═══════╡
│ null │
│ hello │
│ null │
└───────┘
Not turning of predicate_pushdown is better, but then we must ensure the subset is selected before the map.

Answer 2

The answer here provides a technique that can answer this question: gather the columns together in a struct column, and then apply .n_unique() to that struct. That question uses groupby, but it will work without groupby as well.
(
df
.with_column(pl.struct(['a', 'b'].alias('ident'))
['ident'].n_unique()
)
I was able to run code more or less identical to this on a dataset I am working with, and got a sensible answer.
Note that I am not sure if this materializes the entire table before aggregating, nor if this works specifically on lazy data frames. If not, please let me know, and I will retract this answer.

Answer 3

If you have
df=pl.DataFrame({'a':[1,2,3],'b':[2,3,4],'c':[3,4,5]}).lazy()
and you want to see if [a,b] are unique without returning all the data, you can lazily group by and count those groups. With that, you can add a filter such that only rows with a count greater than 1 are returned. Only after those expressions are strung to the LazyFrame do you collect and if your pair of columns remain unique as you intend then the result will have 0 rows.
df \
.groupby(['a','b']) \
.agg(pl.count()). \
filter(pl.col('count')>1).select('count').collect().height