8000 fix(postgres): type a window aggregate over the current row as not null by colinhacks · Pull Request #113 · colinhacks/reylo · GitHub
[go: up one dir, main page]

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 118 additions & 11 deletions engine/postgres/src/infer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,9 @@
//! though, and a small set of SQL constructs carries its own nullability and type:
//!
//! - `count(…)` is never NULL, over no rows and under `GROUP BY` alike. Every other aggregate
//! is NULL over no rows and proves nothing.
//! is NULL over no rows and proves nothing, unless it always has a row: under a plain `GROUP
//! BY`, or as a window function whose frame always holds the current row, it follows its
//! argument.
//! - A literal other than `NULL` is never NULL.
//! - A cast, a `||`, a comparison, `AND`/`OR`/`NOT`, and (on Postgres) arithmetic are non-null
//! when every operand is. SQLite turns a division by zero and a NaN into NULL, so it gets no
Expand Down Expand Up @@ -410,11 +412,15 @@ fn plan_here(sql: &str, dialect: Dialect, trust: &Trust, columns: &Columns) -> P
_ => None,
};
let plain_group_by = matches!(top, NodeEnum::SelectStmt(s) if !grouped && !s.group_clause.is_empty());
let windows: &[Node] = match top {
NodeEnum::SelectStmt(s) => &s.window_clause,
_ => &[],
};
let withs = Withs::default().with(with.as_ref(), columns);
// The old and the new row of a RETURNING, with the relations a bare name is a column of first.
let rows = Rows::of(top).map(|rows| (rows, Level::of(&from, &withs, columns)));
let froms = (dialect == Dialect::Postgres).then(|| vec![Level::of(&from, &withs, columns)]);
let mut reader = Reader { dialect, trust, columns, withs, rows: rows.as_ref(), leaves: (!grouped).then(Vec::new), grouped: plain_group_by, narrowed: Vec::new(), scope: froms };
let mut reader = Reader { dialect, trust, columns, withs, rows: rows.as_ref(), leaves: (!grouped).then(Vec::new), grouped: plain_group_by, narrowed: Vec::new(), scope: froms, windows };
let mut plan = Plan::default();
for target in &targets {
let Some(NodeEnum::ResTarget(rt)) = target.node.as_ref() else { continue };
Expand Down Expand Up @@ -800,7 +806,7 @@ fn array_elements(node: &Node) -> Option<Vec<&Node>> {
/// where the SELECT stands inside another statement: an aggregate counts there only when it is
/// that SELECT's own ([`Level::owns`]).
pub fn aggregates(expr: &Node, dialect: Dialect, trust: &Trust, level: Option<&Level>) -> bool {
Reader { dialect, trust, columns: &Columns::default(), withs: Withs::default(), rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope: None }.aggregates(expr, true, level)
Reader { dialect, trust, columns: &Columns::default(), withs: Withs::default(), rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope: None, windows: &[] }.aggregates(expr, true, level)
}

impl Ty {
Expand Down Expand Up @@ -842,6 +848,8 @@ struct Reader<'a> {
/// them all: what says a bare name in a subquery's aggregate is an item's whole row
/// ([`Level::within`]).
scope: Option<Vec<Level>>,
/// The `WINDOW` clause of the SELECT the reader stands in, which `OVER w` names a window of.
windows: &'a [Node],
}

impl Reader<'_> {
Expand Down Expand Up @@ -1088,10 +1096,11 @@ impl Reader<'_> {
(Dialect::Sqlite, "total", 1) => Expr::Nullable(Box::new(Expr::NonNull(Some(oid::FLOAT8)))),
// "The ifnull() function is equivalent to coalesce() with two arguments."
(Dialect::Sqlite, "ifnull", 2) => Expr::Any(f.args.iter().map(|a| self.read(a)).collect(), Ty::Shared),
// Under a plain GROUP BY every group has a row, so an aggregate that is NULL only over
// no rows or only NULLs follows its argument. A FILTER can leave it no rows, a window
// frame can be empty, and SQLite's sum of `1e999` and `-1e999` is NaN, stored as NULL.
(Dialect::Postgres, n, 1 | 2) if self.grouped && f.over.is_none() && f.agg_filter.is_none() && self.trust.builtin(n) && (POSTGRES_GROUPED.contains(&n) || POSTGRES_GROUPED_NEVER_NULL.contains(&n) || POSTGRES_GROUPED_JSON.contains(&n)) => {
// Under a plain GROUP BY every group has a row, and a window frame that always holds
// the current row has that row, so there an aggregate that is NULL only over no rows
// or only NULLs follows its argument. A FILTER can leave it no rows, any other frame
// can be empty, and SQLite's sum of `1e999` and `-1e999` is NaN, stored as NULL.
(Dialect::Postgres, n, 1 | 2) if f.agg_filter.is_none() && self.has_a_row(f) && self.trust.builtin(n) && (POSTGRES_GROUPED.contains(&n) || POSTGRES_GROUPED_NEVER_NULL.contains(&n) || POSTGRES_GROUPED_JSON.contains(&n)) => {
if POSTGRES_GROUPED_JSON.contains(&n) {
Expr::NonNull(None)
} else if POSTGRES_GROUPED_NEVER_NULL.contains(&n) {
Expand Down Expand Up @@ -1133,6 +1142,16 @@ impl Reader<'_> {
}
}

/// Whether an aggregate call always has a row to aggregate: a plain one under a plain `GROUP
/// BY`, and a window call whose frame always holds the current row
/// ([`frame_holds_current_row`]).
fn has_a_row(&self, f: &protobuf::FuncCall) -> bool {
match f.over.as_deref() {
Some(over) => frame_holds_current_row(over, self.windows),
None => self.grouped,
}
}

/// A read of JSON at literal keys with the built-in `->`, `->>`, `#>` or `#>>`: an
/// [`Expr::JsonPath`] over a column, a read of a read at more keys extending it. A key that
/// is not a literal (a parameter, an expression, an integer that indexes an array), and a read
Expand Down Expand Up @@ -1179,7 +1198,7 @@ impl Reader<'_> {
Some(outside) => level.within(outside.clone()),
None => level,
};
let mut inner_reader = Reader { dialect: self.dialect, trust: self.trust, columns: self.columns, withs, rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope };
let mut inner_reader = Reader { dialect: self.dialect, trust: self.trust, columns: self.columns, withs, rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope, windows: &select.window_clause };
let value = inner_reader.read(inner);
if one_row(select, inner, self.dialect, self.trust) || (keeps_one_row(select, inner, self.dialect, self.trust) && self.aggregates(inner, false, Some(&level))) {
value
Expand Down Expand Up @@ -1802,10 +1821,49 @@ fn nulls_with(node: &Node) -> Vec<Strict> {
}
}

/// Whether the frame of an `OVER` clause always holds the current row ([`holds_current_row`]),
/// given the `WINDOW` clause of the SELECT it stands in. `OVER w` takes the frame of the window
/// `w` that clause defines, and nothing where it defines none; `OVER (w …)` copies `w`'s
/// partitioning and ordering but "always uses its own frame clause" (the server refuses to copy a
/// window that has one), as a window defined `AS (w …)` does.
pub(crate) fn frame_holds_current_row(over: &protobuf::WindowDef, windows: &[Node]) -> bool {
let frame = if over.name.is_empty() {
Some(over.frame_options)
} else {
windows.iter().find_map(|w| match w.node.as_ref() {
Some(NodeEnum::WindowDef(w)) if w.name == over.name => Some(w.frame_options),
_ => None,
})
};
frame.is_some_and(holds_current_row)
}

/// Whether a window frame, as the parser's `frameOptions` bits spell it (`FRAMEOPTION_*` in
/// `parsenodes.h`), holds the current row of every row it is computed for. The default frame,
/// `RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW`, runs through the current row's last peer,
/// and the whole partition without `ORDER BY`. A frame that starts `UNBOUNDED PRECEDING`, `n
/// PRECEDING` or at `CURRENT ROW` and ends at `CURRENT ROW`, `n FOLLOWING` or `UNBOUNDED
/// FOLLOWING` holds it in every mode: an offset "must be non-null and non-negative" (or the
/// statement fails), and in RANGE mode the current row's value is no distance from itself: a
/// NULL is a peer of the NULLs, and a NaN or an infinity, an infinite offset from it included,
/// reaches its own peers (measured on PostgreSQL 18 over float8, numeric and timestamptz). A start
/// `n FOLLOWING` or an end `n PRECEDING` can leave the row out, and `EXCLUDE CURRENT ROW` and
/// `EXCLUDE GROUP` do; `EXCLUDE TIES` keeps it. The two bounds the parser refuses (a start
/// `UNBOUNDED FOLLOWING`, an end `UNBOUNDED PRECEDING`) are refused here too.
fn holds_current_row(options: i32) -> bool {
const START_OFFSET_FOLLOWING: i32 = 0x02000;
const END_OFFSET_PRECEDING: i32 = 0x01000;
const START_UNBOUNDED_FOLLOWING: i32 = 0x00080;
const END_UNBOUNDED_PRECEDING: i32 = 0x00040;
const EXCLUDE_CURRENT_ROW: i32 = 0x08000;
const EXCLUDE_GROUP: i32 = 0x10000;
options & (START_OFFSET_FOLLOWING | END_OFFSET_PRECEDING | START_UNBOUNDED_FOLLOWING | END_UNBOUNDED_PRECEDING | EXCLUDE_CURRENT_ROW | EXCLUDE_GROUP) == 0
}

/// A bare `VALUES (…), (…)` arm: a column is non-null when it is in every row.
fn values_plan(s: &SelectStmt, dialect: Dialect, trust: &Trust, columns: &Columns) -> Plan {
let withs = Withs::default().with(s.with_clause.as_ref(), columns);
let mut reader = Reader { dialect, trust, columns, withs, rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope: None };
let mut reader = Reader { dialect, trust, columns, withs, rows: None, leaves: None, grouped: false, narrowed: Vec::new(), scope: None, windows: &[] };
let mut rows: Vec<Vec<Expr>> = Vec::new();
for row in &s.values_lists {
let Some(NodeEnum::List(list)) = row.node.as_ref() else { return Plan::default() };
Expand Down Expand Up @@ -2281,6 +2339,54 @@ mod tests {
assert!(!one("SELECT first_value(a) OVER () FROM t"));
}

#[test]
fn a_window_aggregate_over_the_current_row_follows_its_argument() {
let &str, non_null: bool| resolve(&pg(sql).head[0].clone().expect("a claim"), Some(&[col(!non_null, oid::INT4)])).non_null;
// No frame clause, with ORDER BY or without, and every frame from at or before the current
// row to at or after it, in every mode.
for over in [
"()",
"(ORDER BY a)",
"(PARTITION BY a)",
"(ORDER BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING)",
"(ORDER BY a ROWS UNBOUNDED PRECEDING)",
"(ORDER BY a ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)",
"(ORDER BY a RANGE BETWEEN 5 PRECEDING AND 0 FOLLOWING)",
"(ORDER BY a GROUPS BETWEEN 2 PRECEDING AND CURRENT ROW)",
"(ORDER BY a RANGE BETWEEN CURRENT ROW AND CURRENT ROW EXCLUDE TIES)",
"(ORDER BY a RANGE UNBOUNDED PRECEDING EXCLUDE NO OTHERS)",
] {
for agg in ["sum", "max", "min", "avg", "bit_or", "array_agg"] {
let sql = format!("SELECT {agg}(a) OVER {over} FROM t");
assert!(one(&sql, true), "{sql}");
// `array_agg` is an array of whatever its rows hold, NULLs included.
assert_eq!(one(&sql, false), agg == "array_agg", "{sql}: a NULL argument");
}
}
assert!(one("SELECT string_agg(a::text, ',') OVER () FROM t", true));
assert!(one("SELECT jsonb_agg(a) OVER (ORDER BY a) FROM t", false), "a JSON array of NULLs is an array");
// A frame that can leave the current row out, and a FILTER, can leave it no row.
for over in [
"(ORDER BY a ROWS BETWEEN 1 FOLLOWING AND 2 FOLLOWING)",
"(ORDER BY a ROWS BETWEEN UNBOUNDED PRECEDING AND 1 PRECEDING)",
"(ORDER BY a ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING EXCLUDE CURRENT ROW)",
"(ORDER BY a RANGE BETWEEN CURRENT ROW AND CURRENT ROW EXCLUDE GROUP)",
] {
let sql = format!("SELECT sum(a) OVER {over} FROM t");
assert!(!one(&sql, true), "{sql}");
}
assert!(!one("SELECT sum(a) FILTER (WHERE a > 1) OVER () FROM t", true));
// A named window lends its frame to `OVER w`; `OVER (w …)` and `AS (w …)` use their own.
assert!(one("SELECT sum(a) OVER w FROM t WINDOW w AS (ORDER BY a)", true));
assert!(!one("SELECT sum(a) OVER w FROM t WINDOW w AS (ORDER BY a ROWS BETWEEN 1 FOLLOWING AND 2 FOLLOWING)", true));
assert!(!one("SELECT sum(a) OVER w FROM t WINDOW v AS (ORDER BY a)", true), "no window w");
assert!(one("SELECT sum(a) OVER (w ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) FROM t WINDOW w AS (PARTITION BY a)", true));
assert!(!one("SELECT sum(a) OVER (w ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING) FROM t WINDOW w AS (PARTITION BY a)", true));
assert!(one("SELECT sum(a) OVER v FROM t WINDOW w AS (PARTITION BY a), v AS (w ORDER BY a ROWS 1 PRECEDING)", true));
// A plain aggregate with no GROUP BY is still NULL over no rows.
assert!(!one("SELECT sum(a) FROM t", true));
}

#[test]
fn sqlite_has_no_total_arithmetic() {
// `1/0` and `1e999 - 1e999` are both NULL on SQLite.
Expand Down Expand Up @@ -2584,14 +2690,15 @@ mod tests {

#[test]
fn a_group_is_never_empty_so_its_aggregate_follows_the_argument() {
let p = pg("SELECT user_id, sum(views) AS total, max(views) AS top, array_agg(title) AS titles, sum(views) FILTER (WHERE views > 0) AS some, sum(views) OVER () AS running FROM posts GROUP BY user_id");
let p = pg("SELECT user_id, sum(views) AS total, max(views) AS top, array_agg(title) AS titles, sum(views) FILTER (WHERE views > 0) AS some, sum(views) OVER (ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) AS next FROM posts GROUP BY user_id");
// One leaf per column reference read: five of `views` here.
let views = [col(false, oid::INT4), col(false, oid::INT4), col(false, oid::INT4), col(false, oid::INT4), col(false, oid::INT4)];
assert_eq!(p.head[0], None);
assert!(resolve(p.head[1].as_ref().unwrap(), Some(&views)).non_null);
assert!(resolve(p.head[2].as_ref().unwrap(), Some(&views)).non_null);
assert!(resolve(p.head[3].as_ref().unwrap(), Some(&views)).non_null);
// A FILTER can leave the group no rows, and a window frame can be empty.
// A FILTER can leave the group no rows, and a frame that leaves out the current row can be
// empty.
assert!(!resolve(p.head[4].as_ref().unwrap(), Some(&views)).non_null);
assert!(!resolve(p.head[5].as_ref().unwrap(), Some(&views)).non_null);
// A nullable argument, no GROUP BY, a grouping set, or SQLite (a sum can be NaN, which it stores as NULL).
Expand Down
23 changes: 17 additions & 6 deletions engine/postgres/src/jsonshape.rs
Original file line number Diff line number Diff line change
Expand Up @@ -176,7 +176,11 @@ fn plan_with(sql: &str, dialect: Dialect, trust: &Trust, columns: &Columns, ahea
NodeEnum::SelectStmt(s) if s.group_clause.is_empty() && aggregating(&targets, dialect, trust) => Some(Vec::new()),
_ => None,
};
let mut builder = Builder { stack: vec![(0, from, filter)], dialect, trust, columns, withs, rows: Rows::of(top), ahead, next_ctx: 0, drafts: Vec::new(), overflowed: false, grouped, group, in_aggregate: 0, grouping_sets, under_aggregate: false, laterals: 0 };
let windows = match top {
NodeEnum::SelectStmt(s) => s.window_clause.clone(),
_ => Vec::new(),
};
let mut builder = Builder { stack: vec![(0, from, filter)], dialect, trust, columns, withs, rows: Rows::of(top), ahead, next_ctx: 0, drafts: Vec::new(), overflowed: false, grouped, group, in_aggregate: 0, grouping_sets, under_aggregate: false, laterals: 0, windows };
let mut columns: Vec<(String, Shape)> = Vec::new();
for target in &targets {
let Some(NodeEnum::ResTarget(rt)) = target.node.as_ref() else { continue };
Expand Down Expand Up @@ -345,6 +349,8 @@ struct Builder<'a> {
under_aggregate: bool,
/// How many LATERAL subqueries' targets have been read, which [`MAX_LATERAL`] bounds.
laterals: usize,
/// The top statement's `WINDOW` clause, which `OVER w` names a window of.
windows: Vec<Node>,
}

/// A target of a LATERAL subquery that builds JSON ([`Builder::lateral_target`]).
Expand Down Expand Up @@ -418,11 +424,16 @@ impl Builder<'_> {
if f.args.len() != 1 {
return Shape::Unknown;
}
// Under a plain GROUP BY every group has a row, so an aggregate of the top
// statement's own is an array there, NULL inputs and all, unless a FILTER or a
// window frame leaves it none. One inside a subquery may be the outer query's
// ([`crate::levels`]), which says nothing of its rows.
let grouped = self.dialect == Dialect::Postgres && self.grouped && self.stack.len() == 1 && f.over.is_none() && f.agg_filter.is_none();
// Under a plain GROUP BY every group has a row, and a window frame that always
// holds the current row has that row, so an aggregate of the top statement's
// own is an array there, NULL inputs and all, unless a FILTER leaves it none.
// One inside a subquery may be the outer query's ([`crate::levels`]), which
// says nothing of its rows.
let has_a_row = match f.over.as_deref() {
Some(over) => crate::infer::frame_holds_current_row(over, &self.windows),
None => self.grouped,
};
let grouped = self.dialect == Dialect::Postgres && has_a_row && self.stack.len() == 1 && f.agg_filter.is_none();
let non_null = TOTAL_AGG_FNS.contains(&name.as_str()) || grouped;
// A plain aggregate of the top statement reads its argument from the rows
// before a grouping set nulls them; a window call reads the grouped rows.
Expand Down
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/aggregate_family.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select avg(p.id) over w as mean, min(p.title) over w as first_title, bool_and(p.id > 0) over w as all_positive, bit_or(p.id) over w as bits, string_agg(p.title, ', ') over w as titles, array_agg(p.title) over w as title_list, array_agg(p.note) over w as note_list, jsonb_agg(p.note) over w as notes, json_object_agg(p.title, p.note) over w as by_title from posts p window w as (partition by p.user_id)
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/around_current.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select sum(p.id) over (order by p.id rows between 1 preceding and 1 following) as rows_around, sum(p.id) over (order by p.id groups between 1 preceding and current row) as groups_back, sum(p.id) over (order by p.id rows between current row and unbounded following) as rest, sum(p.id) over (order by p.id rows unbounded preceding) as running from posts p
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/chained_windows.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select sum(p.id) over w2 as total from posts p window w1 as (partition by p.user_id), w2 as (w1 order by p.id range between unbounded preceding and unbounded following)
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/copied_window.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select sum(p.id) over (w rows between 1 preceding and 1 following) as around, sum(p.id) over (w rows between 2 preceding and 1 preceding) as before from posts p window w as (partition by p.user_id order by p.id)
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/count_window.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select count(*) over () as n from posts p
1 change: 1 addition & 0 deletions oracle/cases/window_aggregates/exclusions.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
select sum(p.id) over (order by p.id rows between unbounded preceding and unbounded following exclude current row) as without_row, sum(p.id) over (order by p.user_id range between current row and current row exclude group) as without_group, sum(p.id) over (order by p.user_id range between current row and current row exclude ties) as without_ties, sum(p.id) over (order by p.id range unbounded preceding exclude no others) as with_all from posts p
Loading
Loading
0