Skip to content

Commit 91fc94f

Browse files
gatesngatesn
committed
Clean up stats rewrite proof helpers
Signed-off-by: "Nicholas Gates" <nick@nickgates.com>
1 parent 9a7c766 commit 91fc94f

1 file changed

Lines changed: 58 additions & 113 deletions

File tree

‎vortex-array/src/stats/rewrite/builtins.rs‎

Lines changed: 58 additions & 113 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ impl StatsRewriteRule for BinaryNanCountStatsRewrite {
8181
expr: &Expression,
8282
ctx: &StatsRewriteCtx<'_>,
8383
) -> VortexResult<Option<Expression>> {
84-
binary_falsify(expr, ctx, nan_count_check, true)
84+
binary_falsify::<NanCountProof>(expr, ctx)
8585
}
8686
}
8787

@@ -98,15 +98,13 @@ impl StatsRewriteRule for BinaryAllNonNanStatsRewrite {
9898
expr: &Expression,
9999
ctx: &StatsRewriteCtx<'_>,
100100
) -> VortexResult<Option<Expression>> {
101-
binary_falsify(expr, ctx, all_non_nan_check, false)
101+
binary_falsify::<AllNonNanProof>(expr, ctx)
102102
}
103103
}
104104

105-
fn binary_falsify(
105+
fn binary_falsify<P: NonNanProof>(
106106
expr: &Expression,
107107
ctx: &StatsRewriteCtx<'_>,
108-
nan_check: fn(&StatsRewriteCtx<'_>, &Expression) -> VortexResult<NanCheck>,
109-
emit_unguarded_rewrites: bool,
110108
) -> VortexResult<Option<Expression>> {
111109
let operator = expr.as_::<Binary>();
112110
let lhs = expr.child(0);
@@ -117,78 +115,44 @@ fn binary_falsify(
117115
let left = min(lhs, ctx).zip(max(rhs, ctx)).map(|(a, b)| gt(a, b));
118116
let right = min(rhs, ctx).zip(max(lhs, ctx)).map(|(a, b)| gt(a, b));
119117
or_collect(left.into_iter().chain(right))
120-
.map(|value_predicate| {
121-
with_nan_predicate(
122-
ctx,
123-
nan_check,
124-
emit_unguarded_rewrites,
125-
lhs,
126-
rhs,
127-
value_predicate,
128-
)
129-
})
118+
.map(|value_predicate| with_non_nan_guards::<P>(ctx, [lhs, rhs], value_predicate))
130119
.transpose()?
131120
.flatten()
132121
}
133122
Operator::NotEq => min(lhs, ctx)
134123
.zip(max(rhs, ctx))
135124
.zip(max(lhs, ctx).zip(min(rhs, ctx)))
136125
.map(|((min_lhs, max_rhs), (max_lhs, min_rhs))| {
137-
with_nan_predicate(
126+
with_non_nan_guards::<P>(
138127
ctx,
139-
nan_check,
140-
emit_unguarded_rewrites,
141-
lhs,
142-
rhs,
128+
[lhs, rhs],
143129
and(eq(min_lhs, max_rhs), eq(max_lhs, min_rhs)),
144130
)
145131
})
146132
.transpose()?
147133
.flatten(),
148134
Operator::Gt => max(lhs, ctx)
149135
.zip(min(rhs, ctx))
150-
.map(|(a, b)| {
151-
with_nan_predicate(
152-
ctx,
153-
nan_check,
154-
emit_unguarded_rewrites,
155-
lhs,
156-
rhs,
157-
lt_eq(a, b),
158-
)
159-
})
136+
.map(|(a, b)| with_non_nan_guards::<P>(ctx, [lhs, rhs], lt_eq(a, b)))
160137
.transpose()?
161138
.flatten(),
162139
Operator::Gte => max(lhs, ctx)
163140
.zip(min(rhs, ctx))
164-
.map(|(a, b)| {
165-
with_nan_predicate(ctx, nan_check, emit_unguarded_rewrites, lhs, rhs, lt(a, b))
166-
})
141+
.map(|(a, b)| with_non_nan_guards::<P>(ctx, [lhs, rhs], lt(a, b)))
167142
.transpose()?
168143
.flatten(),
169144
Operator::Lt => min(lhs, ctx)
170145
.zip(max(rhs, ctx))
171-
.map(|(a, b)| {
172-
with_nan_predicate(
173-
ctx,
174-
nan_check,
175-
emit_unguarded_rewrites,
176-
lhs,
177-
rhs,
178-
gt_eq(a, b),
179-
)
180-
})
146+
.map(|(a, b)| with_non_nan_guards::<P>(ctx, [lhs, rhs], gt_eq(a, b)))
181147
.transpose()?
182148
.flatten(),
183149
Operator::Lte => min(lhs, ctx)
184150
.zip(max(rhs, ctx))
185-
.map(|(a, b)| {
186-
with_nan_predicate(ctx, nan_check, emit_unguarded_rewrites, lhs, rhs, gt(a, b))
187-
})
151+
.map(|(a, b)| with_non_nan_guards::<P>(ctx, [lhs, rhs], gt(a, b)))
188152
.transpose()?
189153
.flatten(),
190154
Operator::And => {
191-
if !emit_unguarded_rewrites {
155+
if !P::EMIT_UNGUARDED_REWRITES {
192156
return Ok(None);
193157
}
194158

@@ -197,7 +161,7 @@ fn binary_falsify(
197161
or_collect(lhs_falsifier.into_iter().chain(rhs_falsifier))
198162
}
199163
Operator::Or => match (ctx.falsify(lhs)?, ctx.falsify(rhs)?) {
200-
(Some(lhs), Some(rhs)) if emit_unguarded_rewrites => Some(and(lhs, rhs)),
164+
(Some(lhs), Some(rhs)) if P::EMIT_UNGUARDED_REWRITES => Some(and(lhs, rhs)),
201165
_ => None,
202166
},
203167
Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None,
@@ -416,7 +380,7 @@ impl StatsRewriteRule for ListContainsNanCountStatsRewrite {
416380
expr: &Expression,
417381
ctx: &StatsRewriteCtx<'_>,
418382
) -> VortexResult<Option<Expression>> {
419-
list_contains_falsify(expr, ctx, nan_count_check, true)
383+
list_contains_falsify::<NanCountProof>(expr, ctx)
420384
}
421385
}
422386

@@ -433,15 +397,13 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite {
433397
expr: &Expression,
434398
ctx: &StatsRewriteCtx<'_>,
435399
) -> VortexResult<Option<Expression>> {
436-
list_contains_falsify(expr, ctx, all_non_nan_check, false)
400+
list_contains_falsify::<AllNonNanProof>(expr, ctx)
437401
}
438402
}
439403

440-
fn list_contains_falsify(
404+
fn list_contains_falsify<P: NonNanProof>(
441405
expr: &Expression,
442406
ctx: &StatsRewriteCtx<'_>,
443-
nan_check: fn(&StatsRewriteCtx<'_>, &Expression) -> VortexResult<NanCheck>,
444-
emit_unguarded_rewrites: bool,
445407
) -> VortexResult<Option<Expression>> {
446408
let list = expr.child(0);
447409
let needle = expr.child(1);
@@ -457,7 +419,7 @@ fn list_contains_falsify(
457419
return Ok(None);
458420
};
459421
if elements.is_empty() {
460-
return Ok(emit_unguarded_rewrites.then(|| lit(true)));
422+
return Ok(P::EMIT_UNGUARDED_REWRITES.then(|| lit(true)));
461423
}
462424

463425
let Some(value_max) = max(needle, ctx) else {
@@ -474,15 +436,7 @@ fn list_contains_falsify(
474436
)
475437
}));
476438
value_predicate
477-
.map(|value_predicate| {
478-
with_all_non_nan_predicate(
479-
ctx,
480-
nan_check,
481-
emit_unguarded_rewrites,
482-
[needle],
483-
value_predicate,
484-
)
485-
})
439+
.map(|value_predicate| with_non_nan_guards::<P>(ctx, [needle], value_predicate))
486440
.transpose()
487441
.map(Option::flatten)
488442
}
@@ -500,7 +454,7 @@ impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite {
500454
expr: &Expression,
501455
ctx: &StatsRewriteCtx<'_>,
502456
) -> VortexResult<Option<Expression>> {
503-
dynamic_comparison_falsify(expr, ctx, nan_count_check, true)
457+
dynamic_comparison_falsify::<NanCountProof>(expr, ctx)
504458
}
505459
}
506460

@@ -517,15 +471,13 @@ impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite {
517471
expr: &Expression,
518472
ctx: &StatsRewriteCtx<'_>,
519473
) -> VortexResult<Option<Expression>> {
520-
dynamic_comparison_falsify(expr, ctx, all_non_nan_check, false)
474+
dynamic_comparison_falsify::<AllNonNanProof>(expr, ctx)
521475
}
522476
}
523477

524-
fn dynamic_comparison_falsify(
478+
fn dynamic_comparison_falsify<P: NonNanProof>(
525479
expr: &Expression,
526480
ctx: &StatsRewriteCtx<'_>,
527-
nan_check: fn(&StatsRewriteCtx<'_>, &Expression) -> VortexResult<NanCheck>,
528-
emit_unguarded_rewrites: bool,
529481
) -> VortexResult<Option<Expression>> {
530482
let dynamic = expr.as_::<DynamicComparison>();
531483
let lhs = expr.child(0);
@@ -548,13 +500,7 @@ fn dynamic_comparison_falsify(
548500
},
549501
[lhs_stat],
550502
);
551-
with_all_non_nan_predicate(
552-
ctx,
553-
nan_check,
554-
emit_unguarded_rewrites,
555-
[lhs],
556-
value_predicate,
557-
)
503+
with_non_nan_guards::<P>(ctx, [lhs], value_predicate)
558504
}
559505

560506
fn min(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option<Expression> {
@@ -583,25 +529,43 @@ enum NanCheck {
583529
Unavailable,
584530
}
585531

586-
// Min/max do not order NaN values, so comparison rewrites are only sound when every
587-
// candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast
588-
// from float to non-float still needs a proof about the float source values.
589-
fn nan_count_check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult<NanCheck> {
590-
nan_check(ctx, expr, |expr| {
591-
match stat_expr(expr, Stat::NaNCount, ctx) {
592-
Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))),
593-
None => NanCheck::Unavailable,
594-
}
595-
})
532+
trait NonNanProof {
533+
const EMIT_UNGUARDED_REWRITES: bool;
534+
535+
fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult<NanCheck>;
596536
}
597537

598-
fn all_non_nan_check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult<NanCheck> {
599-
nan_check(ctx, expr, |expr| {
600-
NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)))
601-
})
538+
struct NanCountProof;
539+
540+
impl NonNanProof for NanCountProof {
541+
const EMIT_UNGUARDED_REWRITES: bool = true;
542+
543+
fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult<NanCheck> {
544+
non_nan_check(ctx, expr, |expr| {
545+
match stat_expr(expr, Stat::NaNCount, ctx) {
546+
Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))),
547+
None => NanCheck::Unavailable,
548+
}
549+
})
550+
}
551+
}
552+
553+
struct AllNonNanProof;
554+
555+
impl NonNanProof for AllNonNanProof {
556+
const EMIT_UNGUARDED_REWRITES: bool = false;
557+
558+
fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult<NanCheck> {
559+
non_nan_check(ctx, expr, |expr| {
560+
NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions)))
561+
})
562+
}
602563
}
603564

604-
fn nan_check(
565+
// Min/max do not order NaN values, so comparison rewrites are only sound when every
566+
// candidate value is known to be non-NaN. Cast result dtypes are not enough: a cast
567+
// from float to non-float still needs a proof about the float source values.
568+
fn non_nan_check(
605569
ctx: &StatsRewriteCtx<'_>,
606570
expr: &Expression,
607571
proof: impl FnOnce(&Expression) -> NanCheck,
@@ -622,7 +586,7 @@ fn nan_check(
622586
return Ok(NanCheck::NotNeeded);
623587
}
624588

625-
return nan_check(ctx, expr.child(0), proof);
589+
return non_nan_check(ctx, expr.child(0), proof);
626590
}
627591

628592
if !has_nans(&ctx.return_dtype(expr)?) {
@@ -663,33 +627,14 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option
663627
.then(|| stat_fn(expr.clone(), aggregate_fn))
664628
}
665629

666-
fn with_nan_predicate(
667-
ctx: &StatsRewriteCtx<'_>,
668-
nan_check: fn(&StatsRewriteCtx<'_>, &Expression) -> VortexResult<NanCheck>,
669-
emit_unguarded_rewrites: bool,
670-
lhs: &Expression,
671-
rhs: &Expression,
672-
value_predicate: Expression,
673-
) -> VortexResult<Option<Expression>> {
674-
with_all_non_nan_predicate(
675-
ctx,
676-
nan_check,
677-
emit_unguarded_rewrites,
678-
[lhs, rhs],
679-
value_predicate,
680-
)
681-
}
682-
683-
fn with_all_non_nan_predicate<'a>(
630+
fn with_non_nan_guards<'a, P: NonNanProof>(
684631
ctx: &StatsRewriteCtx<'_>,
685-
nan_check: fn(&StatsRewriteCtx<'_>, &Expression) -> VortexResult<NanCheck>,
686-
emit_unguarded_rewrites: bool,
687632
exprs: impl IntoIterator<Item = &'a Expression>,
688633
value_predicate: Expression,
689634
) -> VortexResult<Option<Expression>> {
690635
let mut nan_checks = Vec::new();
691636
for expr in exprs {
692-
match nan_check(ctx, expr)? {
637+
match P::check(ctx, expr)? {
693638
NanCheck::NotNeeded => {}
694639
NanCheck::Check(check) => nan_checks.push(check),
695640
NanCheck::Unavailable => return Ok(None),
@@ -702,7 +647,7 @@ fn with_all_non_nan_predicate<'a>(
702647
// No possible NaN-bearing expression remains, so the value predicate is
703648
// already guarded. Only one registered rule emits this unguarded
704649
// rewrite so non-float comparisons are not duplicated.
705-
None if emit_unguarded_rewrites => Some(value_predicate),
650+
None if P::EMIT_UNGUARDED_REWRITES => Some(value_predicate),
706651
None => None,
707652
})
708653
}

0 commit comments

Comments
 (0)