Skip to content

Commit

Permalink
[Postgres] Add order set aggregate functions support.
Browse files Browse the repository at this point in the history
  • Loading branch information
go-jet committed Jan 18, 2022
1 parent 04c14f2 commit 605f1c8
Show file tree
Hide file tree
Showing 4 changed files with 135 additions and 1 deletion.
60 changes: 60 additions & 0 deletions internal/jet/order_set_aggregate_functions.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
package jet

// MODE computes the most frequent value of the aggregated argument
func MODE() *OrderSetAggregateFunc {
return newOrderSetAggregateFunction("MODE", nil)
}

// PERCENTILE_CONT computes a value corresponding to the specified fraction within the ordered set of
// aggregated argument values. This will interpolate between adjacent input items if needed.
func PERCENTILE_CONT(fraction FloatExpression) *OrderSetAggregateFunc {
return newOrderSetAggregateFunction("PERCENTILE_CONT", fraction)
}

// PERCENTILE_DISC computes the first value within the ordered set of aggregated argument values whose position
// in the ordering equals or exceeds the specified fraction. The aggregated argument must be of a sortable type.
func PERCENTILE_DISC(fraction FloatExpression) *OrderSetAggregateFunc {
return newOrderSetAggregateFunction("PERCENTILE_DISC", fraction)
}

// OrderSetAggregateFunc implementation of order set aggregate function
type OrderSetAggregateFunc struct {
name string
fraction FloatExpression
orderBy Window
}

func newOrderSetAggregateFunction(name string, fraction FloatExpression) *OrderSetAggregateFunc {
return &OrderSetAggregateFunc{
name: name,
fraction: fraction,
}
}

// WITHIN_GROUP_ORDER_BY specifies ordered set of aggregated argument values
func (p *OrderSetAggregateFunc) WITHIN_GROUP_ORDER_BY(orderBy OrderByClause) Expression {
p.orderBy = ORDER_BY(orderBy)
return newOrderSetAggregateFuncExpression(*p)
}

func newOrderSetAggregateFuncExpression(aggFunc OrderSetAggregateFunc) *orderSetAggregateFuncExpression {
ret := &orderSetAggregateFuncExpression{
OrderSetAggregateFunc: aggFunc,
}

ret.ExpressionInterfaceImpl.Parent = ret

return ret
}

type orderSetAggregateFuncExpression struct {
ExpressionInterfaceImpl
OrderSetAggregateFunc
}

func (p *orderSetAggregateFuncExpression) serialize(statement StatementType, out *SQLBuilder, options ...SerializeOption) {
out.WriteString(p.name)
WRAP(p.fraction).serialize(statement, out, FallTrough(options)...)
out.WriteString("WITHIN GROUP")
p.orderBy.serialize(statement, out)
}
4 changes: 3 additions & 1 deletion internal/jet/utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,9 @@ func serializeExpressionList(
out.WriteString(separator)
}

expression.serialize(statement, out, options...)
if expression != nil {
expression.serialize(statement, out, options...)
}
}
}

Expand Down
22 changes: 22 additions & 0 deletions postgres/functions.go
Original file line number Diff line number Diff line change
Expand Up @@ -336,3 +336,25 @@ func explicitLiteralCast(expresion Expression) jet.Expression {

return expresion
}

// MODE computes the most frequent value of the aggregated argument
var MODE = jet.MODE

// PERCENTILE_CONT computes a value corresponding to the specified fraction within the ordered set of
// aggregated argument values. This will interpolate between adjacent input items if needed.
func PERCENTILE_CONT(fraction FloatExpression) *jet.OrderSetAggregateFunc {
return jet.PERCENTILE_CONT(castFloatLiteral(fraction))
}

// PERCENTILE_DISC computes the first value within the ordered set of aggregated argument values whose position
// in the ordering equals or exceeds the specified fraction. The aggregated argument must be of a sortable type.
func PERCENTILE_DISC(fraction FloatExpression) *jet.OrderSetAggregateFunc {
return jet.PERCENTILE_DISC(castFloatLiteral(fraction))
}

func castFloatLiteral(fraction FloatExpression) FloatExpression {
if _, ok := fraction.(jet.LiteralExpression); ok {
return CAST(fraction).AS_DOUBLE() // to make postgres aware of the type
}
return fraction
}
50 changes: 50 additions & 0 deletions tests/postgres/chinook_db_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -743,3 +743,53 @@ var album347 = model.Album{
Title: "Koyaanisqatsi (Soundtrack from the Motion Picture)",
ArtistId: 275,
}

func TestAggregateFunc(t *testing.T) {
stmt := SELECT(
PERCENTILE_DISC(Float(0.1)).WITHIN_GROUP_ORDER_BY(Invoice.InvoiceId).AS("percentile_disc_1"),
PERCENTILE_DISC(Invoice.Total.DIV(Float(100))).WITHIN_GROUP_ORDER_BY(Invoice.InvoiceDate.ASC()).AS("percentile_disc_2"),
PERCENTILE_DISC(RawFloat("(select array_agg(s) from generate_series(0, 1, 0.2) as s)")).
WITHIN_GROUP_ORDER_BY(Invoice.BillingAddress.DESC()).AS("percentile_disc_3"),

PERCENTILE_CONT(Float(0.3)).WITHIN_GROUP_ORDER_BY(Invoice.Total).AS("percentile_cont_1"),
PERCENTILE_CONT(Float(0.2)).WITHIN_GROUP_ORDER_BY(INTERVAL(1, HOUR).DESC()).AS("percentile_cont_int"),

MODE().WITHIN_GROUP_ORDER_BY(Invoice.BillingPostalCode.DESC()).AS("mode_1"),
).FROM(
Invoice,
).GROUP_BY(
Invoice.Total,
)

testutils.AssertStatementSql(t, stmt, `
SELECT PERCENTILE_DISC ($1::double precision) WITHIN GROUP (ORDER BY "Invoice"."InvoiceId") AS "percentile_disc_1",
PERCENTILE_DISC ("Invoice"."Total" / $2) WITHIN GROUP (ORDER BY "Invoice"."InvoiceDate" ASC) AS "percentile_disc_2",
PERCENTILE_DISC ((select array_agg(s) from generate_series(0, 1, 0.2) as s)) WITHIN GROUP (ORDER BY "Invoice"."BillingAddress" DESC) AS "percentile_disc_3",
PERCENTILE_CONT ($3::double precision) WITHIN GROUP (ORDER BY "Invoice"."Total") AS "percentile_cont_1",
PERCENTILE_CONT ($4::double precision) WITHIN GROUP (ORDER BY INTERVAL '1 HOUR' DESC) AS "percentile_cont_int",
MODE () WITHIN GROUP (ORDER BY "Invoice"."BillingPostalCode" DESC) AS "mode_1"
FROM chinook."Invoice"
GROUP BY "Invoice"."Total";
`, 0.1, 100.0, 0.3, 0.2)

var dest struct {
PercentileDisc1 string
PercentileDisc2 string
PercentileDisc3 string
PercentileCont1 string
Mode1 string
}

err := stmt.Query(db, &dest)

require.NoError(t, err)
testutils.AssertJSON(t, dest, `
{
"PercentileDisc1": "41",
"PercentileDisc2": "2009-01-19T00:00:00Z",
"PercentileDisc3": "{\"Via Degli Scipioni, 43\",\"Qe 7 Bloco G\",\"Berger Stra�e 10\",\"696 Osborne Street\",\"2211 W Berry Street\",\"1033 N Park Ave\"}",
"PercentileCont1": "0.99",
"Mode1": "X1A 1N6"
}
`)
}

0 comments on commit 605f1c8

Please sign in to comment.