Skip to content

Commit 32241d2

Browse files
gatesngatesn
committed
Make legacy stat aggregate binding explicit
Signed-off-by: "Nicholas Gates" <nick@nickgates.com>
1 parent 118c995 commit 32241d2

4 files changed

Lines changed: 161 additions & 35 deletions

File tree

‎vortex-array/src/expr/pruning/pruning_expr.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use vortex_session::VortexSession;
99
use vortex_utils::aliases::hash_set::HashSet;
1010

1111
use super::relation::Relation;
12+
use crate::aggregate_fn::AggregateFnRef;
1213
use crate::dtype::DType;
1314
use crate::dtype::Field;
1415
use crate::dtype::FieldName;
@@ -26,6 +27,7 @@ use crate::scalar_fn::fns::cast::Cast;
2627
use crate::scalar_fn::fns::get_item::GetItem;
2728
use crate::scalar_fn::fns::literal::Literal;
2829
use crate::stats::bind::StatBinder;
30+
use crate::stats::bind::bind_legacy_count_or_direct_aggregate;
2931
use crate::stats::bind::bind_stats;
3032

3133
pub type RequiredStats = Relation<FieldPath, Stat>;
@@ -139,6 +141,15 @@ impl StatBinder for RequiredStatsBinder<'_> {
139141
self.required_stats.insert(field_path, stat);
140142
Ok(Some(get_item(stat_field_name, root())))
141143
}
144+
145+
fn bind_aggregate(
146+
&mut self,
147+
input: &Expression,
148+
aggregate_fn: &AggregateFnRef,
149+
stat_dtype: &DType,
150+
) -> VortexResult<Option<Expression>> {
151+
bind_legacy_count_or_direct_aggregate(self, input, aggregate_fn, stat_dtype)
152+
}
142153
}
143154

144155
fn direct_stat_field_path(expr: &Expression) -> Option<FieldPath> {

‎vortex-array/src/stats/bind.rs‎

Lines changed: 128 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -64,41 +64,7 @@ pub trait StatBinder {
6464
aggregate_fn: &AggregateFnRef,
6565
stat_dtype: &DType,
6666
) -> VortexResult<Option<Expression>> {
67-
// These aggregate stats are derived from legacy count slots rather than
68-
// stored directly. Keep that storage mapping in binding so stats rewrite
69-
// rules can continue to ask for the semantic aggregate they need.
70-
if aggregate_fn.is::<AllNan>() {
71-
let Some(nan_count) = self.bind_legacy_stat(input, Stat::NaNCount)? else {
72-
return Ok(None);
73-
};
74-
return Ok(Some(eq(nan_count, RowCount.new_expr(EmptyOptions, []))));
75-
}
76-
77-
if aggregate_fn.is::<AllNonNan>() {
78-
let Some(nan_count) = self.bind_legacy_stat(input, Stat::NaNCount)? else {
79-
return Ok(None);
80-
};
81-
return Ok(Some(eq(nan_count, lit(0u64))));
82-
}
83-
84-
if aggregate_fn.is::<AllNull>() {
85-
let Some(null_count) = self.bind_legacy_stat(input, Stat::NullCount)? else {
86-
return Ok(None);
87-
};
88-
return Ok(Some(eq(null_count, RowCount.new_expr(EmptyOptions, []))));
89-
}
90-
91-
if aggregate_fn.is::<AllNonNull>() {
92-
let Some(null_count) = self.bind_legacy_stat(input, Stat::NullCount)? else {
93-
return Ok(None);
94-
};
95-
return Ok(Some(eq(null_count, lit(0u64))));
96-
}
97-
98-
let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else {
99-
return Ok(None);
100-
};
101-
self.bind_stat(input, stat, stat_dtype)
67+
bind_direct_aggregate_stat(self, input, aggregate_fn, stat_dtype)
10268
}
10369

10470
/// Bind one of the legacy stat slots for `input`.
@@ -149,6 +115,77 @@ pub fn bind_stats(predicate: Expression, binder: &mut impl StatBinder) -> Vortex
149115
lowered.optimize_recursive(&binder.bound_scope())
150116
}
151117

118+
/// Bind aggregate stats that can be derived from legacy count stat slots.
119+
///
120+
/// This is an opt-in helper for stats backends that materialize `NaNCount` and
121+
/// `NullCount`, but do not materialize aggregate boolean stats directly.
122+
pub fn bind_legacy_count_aggregate<B: StatBinder + ?Sized>(
123+
binder: &mut B,
124+
input: &Expression,
125+
aggregate_fn: &AggregateFnRef,
126+
) -> VortexResult<Option<Expression>> {
127+
if aggregate_fn.is::<AllNan>() {
128+
let Some(nan_count) = binder.bind_legacy_stat(input, Stat::NaNCount)? else {
129+
return Ok(None);
130+
};
131+
return Ok(Some(eq(nan_count, RowCount.new_expr(EmptyOptions, []))));
132+
}
133+
134+
if aggregate_fn.is::<AllNonNan>() {
135+
let Some(nan_count) = binder.bind_legacy_stat(input, Stat::NaNCount)? else {
136+
return Ok(None);
137+
};
138+
return Ok(Some(eq(nan_count, lit(0u64))));
139+
}
140+
141+
if aggregate_fn.is::<AllNull>() {
142+
let Some(null_count) = binder.bind_legacy_stat(input, Stat::NullCount)? else {
143+
return Ok(None);
144+
};
145+
return Ok(Some(eq(null_count, RowCount.new_expr(EmptyOptions, []))));
146+
}
147+
148+
if aggregate_fn.is::<AllNonNull>() {
149+
let Some(null_count) = binder.bind_legacy_stat(input, Stat::NullCount)? else {
150+
return Ok(None);
151+
};
152+
return Ok(Some(eq(null_count, lit(0u64))));
153+
}
154+
155+
Ok(None)
156+
}
157+
158+
/// Bind an aggregate function that has a direct legacy [`Stat`] slot.
159+
pub fn bind_direct_aggregate_stat<B: StatBinder + ?Sized>(
160+
binder: &mut B,
161+
input: &Expression,
162+
aggregate_fn: &AggregateFnRef,
163+
stat_dtype: &DType,
164+
) -> VortexResult<Option<Expression>> {
165+
let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else {
166+
return Ok(None);
167+
};
168+
binder.bind_stat(input, stat, stat_dtype)
169+
}
170+
171+
/// Bind aggregate stats for backends that expose legacy count-derived stats.
172+
///
173+
/// Backends using this helper first bind aggregate facts derivable from
174+
/// `NaNCount` and `NullCount`, then fall back to direct aggregate-to-stat
175+
/// mappings.
176+
pub fn bind_legacy_count_or_direct_aggregate<B: StatBinder + ?Sized>(
177+
binder: &mut B,
178+
input: &Expression,
179+
aggregate_fn: &AggregateFnRef,
180+
stat_dtype: &DType,
181+
) -> VortexResult<Option<Expression>> {
182+
if let Some(bound) = bind_legacy_count_aggregate(binder, input, aggregate_fn)? {
183+
return Ok(Some(bound));
184+
}
185+
186+
bind_direct_aggregate_stat(binder, input, aggregate_fn, stat_dtype)
187+
}
188+
152189
fn bind_stat_fn(
153190
expr: &Expression,
154191
scope: &DType,
@@ -175,9 +212,11 @@ mod tests {
175212
use crate::dtype::Nullability;
176213
use crate::dtype::PType;
177214
use crate::dtype::StructFields;
215+
use crate::expr::and;
178216
use crate::expr::col;
179217
use crate::expr::get_item;
180218
use crate::expr::is_null;
219+
use crate::expr::or;
181220
use crate::expr::root;
182221
use crate::stats::all_non_nan;
183222

@@ -230,6 +269,15 @@ mod tests {
230269
Ok(None)
231270
}
232271
}
272+
273+
fn bind_aggregate(
274+
&mut self,
275+
input: &Expression,
276+
aggregate_fn: &AggregateFnRef,
277+
stat_dtype: &DType,
278+
) -> VortexResult<Option<Expression>> {
279+
bind_legacy_count_or_direct_aggregate(self, input, aggregate_fn, stat_dtype)
280+
}
233281
}
234282

235283
#[test]
@@ -252,6 +300,51 @@ mod tests {
252300
Ok(())
253301
}
254302

303+
#[test]
304+
fn missing_stats_fold_when_kleene_semantics_allow_it() -> VortexResult<()> {
305+
let mut binder = TestBinder::new(false);
306+
307+
let bound = bind_stats(and(lit(false), all_non_nan(col("f"))), &mut binder)?;
308+
309+
assert_eq!(bound, lit(false));
310+
311+
let bound = bind_stats(or(lit(true), all_non_nan(col("f"))), &mut binder)?;
312+
313+
assert_eq!(bound, lit(true));
314+
Ok(())
315+
}
316+
317+
#[test]
318+
fn default_binder_does_not_derive_all_non_nan_from_nan_count() -> VortexResult<()> {
319+
struct DefaultBinder(TestBinder);
320+
321+
impl StatBinder for DefaultBinder {
322+
fn scope(&self) -> &DType {
323+
self.0.scope()
324+
}
325+
326+
fn bound_scope(&self) -> DType {
327+
self.0.bound_scope()
328+
}
329+
330+
fn bind_stat(
331+
&mut self,
332+
input: &Expression,
333+
stat: Stat,
334+
stat_dtype: &DType,
335+
) -> VortexResult<Option<Expression>> {
336+
self.0.bind_stat(input, stat, stat_dtype)
337+
}
338+
}
339+
340+
let mut binder = DefaultBinder(TestBinder::new(true));
341+
342+
let bound = bind_stats(all_non_nan(col("f")), &mut binder)?;
343+
344+
assert_eq!(bound, lit(Scalar::null(DType::Bool(Nullability::Nullable))));
345+
Ok(())
346+
}
347+
255348
#[test]
256349
fn unrelated_expressions_do_not_request_nan_count() -> VortexResult<()> {
257350
let mut binder = TestBinder::new(false);

‎vortex-file/src/v2/file_stats_reader.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use vortex_array::Canonical;
1414
use vortex_array::IntoArray;
1515
use vortex_array::MaskFuture;
1616
use vortex_array::VortexSessionExecute;
17+
use vortex_array::aggregate_fn::AggregateFnRef;
1718
use vortex_array::arrays::ConstantArray;
1819
use vortex_array::arrays::NullArray;
1920
use vortex_array::dtype::DType;
@@ -29,6 +30,7 @@ use vortex_array::scalar_fn::fns::get_item::GetItem;
2930
use vortex_array::scalar_fn::fns::literal::Literal;
3031
use vortex_array::scalar_fn::internal::row_count::substitute_row_count;
3132
use vortex_array::stats::bind::StatBinder;
33+
use vortex_array::stats::bind::bind_legacy_count_or_direct_aggregate;
3234
use vortex_array::stats::bind::bind_stats;
3335
use vortex_error::VortexResult;
3436
use vortex_layout::ArrayFuture;
@@ -152,6 +154,15 @@ impl StatBinder for FileStatsBinder<'_> {
152154
};
153155
Ok(self.reader.stat_ref(&field_path, stat))
154156
}
157+
158+
fn bind_aggregate(
159+
&mut self,
160+
input: &Expression,
161+
aggregate_fn: &AggregateFnRef,
162+
stat_dtype: &DType,
163+
) -> VortexResult<Option<Expression>> {
164+
bind_legacy_count_or_direct_aggregate(self, input, aggregate_fn, stat_dtype)
165+
}
155166
}
156167

157168
fn direct_field_path(expr: &Expression) -> Option<FieldPath> {

‎vortex-layout/src/layouts/zoned/zone_map.rs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use std::sync::Arc;
88
use vortex_array::ArrayRef;
99
use vortex_array::IntoArray;
1010
use vortex_array::VortexSessionExecute;
11+
use vortex_array::aggregate_fn::AggregateFnRef;
1112
use vortex_array::arrays::ConstantArray;
1213
use vortex_array::arrays::PrimitiveArray;
1314
use vortex_array::arrays::StructArray;
@@ -22,6 +23,7 @@ use vortex_array::expr::stats::Stat;
2223
use vortex_array::scalar_fn::internal::row_count::contains_row_count;
2324
use vortex_array::scalar_fn::internal::row_count::substitute_row_count;
2425
use vortex_array::stats::bind::StatBinder;
26+
use vortex_array::stats::bind::bind_legacy_count_or_direct_aggregate;
2527
use vortex_array::stats::bind::bind_stats;
2628
use vortex_array::validity::Validity;
2729
use vortex_buffer::buffer;
@@ -157,6 +159,15 @@ impl StatBinder for ZoneMapStatsBinder<'_> {
157159
}
158160
Ok(Some(get_item(stat.name(), root())))
159161
}
162+
163+
fn bind_aggregate(
164+
&mut self,
165+
input: &Expression,
166+
aggregate_fn: &AggregateFnRef,
167+
stat_dtype: &DType,
168+
) -> VortexResult<Option<Expression>> {
169+
bind_legacy_count_or_direct_aggregate(self, input, aggregate_fn, stat_dtype)
170+
}
160171
}
161172

162173
/// Build per-zone row counts for a zone map.

0 commit comments

Comments
 (0)