Aggregate function filtering

[oeed]

I've found myself needing Postgres' aggregate function filtering, e.g.

SELECT COUNT(*), COUNT(*) FILTER (WHERE "first_name" = 'Alice') FROM "user";

It appears to not be implemented yet, so I've sketched out this rough implementation if you're open to a PR for it?

/// Marker trait implemented on all aggregate functions (would add to the define_sql_function macro)
pub trait AggregateFunction {}

pub trait AggregateFilterDsl<Predicate> {
  /// The type returned by `.filter`.
  type Output;

  /// See the trait documentation.
  fn filter(self, predicate: Predicate) -> Self::Output;
}

impl AggregateFunction for CountStar {}
impl<T, E> AggregateFunction for CountDistinct<T, E> {}
impl<T: SqlType + SingleValue, Expr> AggregateFunction for count<T, Expr> where Expr: Expression {}
impl<T: Foldable, Expr> AggregateFunction for sum<T, Expr> where Expr: Expression {}

impl<T, Predicate> AggregateFilterDsl<Predicate> for T
where
  T: AggregateFunction,
  Predicate: Expression,
  Predicate::SqlType: BoolOrNullableBool,
{
  type Output = AggregateFilter<T, WhereClause<Predicate>>;

  fn filter(self, predicate: Predicate) -> Self::Output {
    AggregateFilter {
      agg: self,
      where_clause: NoWhereClause.and(predicate),
    }
  }
}

impl<T, Predicate, W> AggregateFilterDsl<Predicate> for AggregateFilter<T, WhereClause<W>>
where
  T: AggregateFunction,
  Predicate: Expression,
  Predicate::SqlType: BoolOrNullableBool,
  W: Expression,
  W::SqlType: BoolOrNullableBool,
{
  type Output = AggregateFilter<T, <WhereClause<W> as WhereAnd<Predicate>>::Output>;

  fn filter(self, predicate: Predicate) -> Self::Output {
    AggregateFilter {
      agg: self.agg,
      where_clause: self.where_clause.and(predicate),
    }
  }
}

pub struct AggregateFilter<A, W> {
  agg: A,
  where_clause: W,
}

impl<A, W> Expression for AggregateFilter<A, W>
where
  A: Expression,
{
  type SqlType = A::SqlType;
}

impl<A, Predicate> QueryFragment<Pg> for AggregateFilter<A, Predicate>
where
  A: QueryFragment<Pg>,
  Predicate: QueryFragment<Pg>,
{
  fn walk_ast<'b>(&'b self, mut out: AstPass<'_, 'b, Pg>) -> QueryResult<()> {
    self.agg.walk_ast(out.reborrow())?;
    out.push_sql(" FILTER (");
    self.where_clause.walk_ast(out.reborrow())?;
    out.push_sql(")");
    Ok(())
  }
}

Usage:

user::table
  .select((
    count_star(),
    count_star().filter(user_schema::first_name.eq("Alice")),
  ))
  .get_result(conn);

[weiznich]

Thanks for writing this proposal. On it's own the proposed extension looks reasonable, but I fear we do need to perform more design work here as the postgres documentation outlines quite a few other syntax variants there that might interact with filter clauses in this location. This becomes especially important if you also want to support window functions which have a syntax part that looks essentially the same. We also need to research if other backends might support this syntax as well. At least for window functions (+filter) that seems the case for SQLite.

The good news is that I'm already working on implementing window functions so there will likely be a better starting point published in the next few weeks. At least in my local draft this also enables aggregate(DISTINCT column) support for all aggregates and so on. I would be happy to receive a review of the proposed API as soon as the changeset is pushed.

[oeed]

Okay great thanks, feel free to @ me once ready