-
Notifications
You must be signed in to change notification settings - Fork 203
Expand file tree
/
Copy pathrules.rs
More file actions
228 lines (207 loc) · 8.14 KB
/
Copy pathrules.rs
File metadata and controls
228 lines (207 loc) · 8.14 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors
//! Metadata-only rewrite rules for the optimizer (Layers 1 and 2 of the execution model).
//!
//! Reduce rules are the cheapest transformations in the execution pipeline: they operate
//! purely on array structure and metadata without reading any data buffers.
//!
//! There are two kinds of reduce rules:
//!
//! - [`ArrayReduceRule`] (Layer 1) -- a self-rewrite where an array simplifies itself.
//! Example: a `FilterArray` with an all-true mask removes the filter wrapper.
//!
//! - [`ArrayParentReduceRule`] (Layer 2) -- a child-driven rewrite where a child rewrites
//! its parent. Example: a `DictArray` child of a `ScalarFnArray` pushes the scalar function
//! into the dictionary values.
//!
//! Rules are collected into [`ReduceRuleSet`] and [`ParentRuleSet`] respectively, and
//! evaluated by the optimizer in a fixpoint loop until no more rules apply.
use std::any::type_name;
use std::fmt::Debug;
use std::marker::PhantomData;
use vortex_error::VortexResult;
use crate::ArrayRef;
use crate::array::ArrayView;
use crate::array::VTable;
use crate::matcher::Matcher;
use crate::trace_op;
/// A metadata-only rewrite rule that transforms an array based on its own structure (Layer 1).
///
/// These rules look only at the array's metadata and children types (not buffer contents)
/// and return a structurally simpler replacement, or `None` if the rule doesn't apply.
pub trait ArrayReduceRule<V: VTable>: Debug + Send + Sync + 'static {
/// Attempt to rewrite this array.
///
/// Returns:
/// - `Ok(Some(new_array))` if the rule applied successfully
/// - `Ok(None)` if the rule doesn't apply
/// - `Err(e)` if an error occurred
fn reduce(&self, array: ArrayView<'_, V>) -> VortexResult<Option<ArrayRef>>;
}
/// A metadata-only rewrite rule where a child encoding rewrites its parent (Layer 2).
///
/// The child sees the parent's type via the associated `Parent` [`Matcher`] and can return
/// a replacement for the parent. This enables optimizations like pushing operations through
/// compression layers (e.g., pushing a scalar function into dictionary values).
pub trait ArrayParentReduceRule<V: VTable>: Debug + Send + Sync + 'static {
/// The parent array type this rule matches against.
type Parent: Matcher;
/// Attempt to rewrite this child array given information about its parent.
///
/// Returns:
/// - `Ok(Some(new_array))` if the rule applied successfully
/// - `Ok(None)` if the rule doesn't apply
/// - `Err(e)` if an error occurred
fn reduce_parent(
&self,
array: ArrayView<'_, V>,
parent: <Self::Parent as Matcher>::Match<'_>,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>>;
}
/// Type-erased version of [`ArrayParentReduceRule`] used for dynamic dispatch within
/// [`ParentRuleSet`].
pub trait DynArrayParentReduceRule<V: VTable>: Debug + Send + Sync {
fn matches(&self, parent: &ArrayRef) -> bool;
fn reduce_parent(
&self,
array: ArrayView<'_, V>,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>>;
}
/// Bridges a concrete [`ArrayParentReduceRule<V, R>`] to the type-erased
/// [`DynArrayParentReduceRule<V>`] trait. Created by [`ParentRuleSet::lift`].
pub struct ParentReduceRuleAdapter<V, R> {
rule: R,
_phantom: PhantomData<V>,
}
impl<V: VTable, R: ArrayParentReduceRule<V>> Debug for ParentReduceRuleAdapter<V, R> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ArrayParentReduceRuleAdapter")
.field("parent", &type_name::<R::Parent>())
.field("rule", &self.rule)
.finish()
}
}
impl<V: VTable, K: ArrayParentReduceRule<V>> DynArrayParentReduceRule<V>
for ParentReduceRuleAdapter<V, K>
{
fn matches(&self, parent: &ArrayRef) -> bool {
K::Parent::matches(parent)
}
fn reduce_parent(
&self,
child: ArrayView<'_, V>,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>> {
let Some(parent_view) = K::Parent::try_match(parent) else {
return Ok(None);
};
self.rule.reduce_parent(child, parent_view, child_idx)
}
}
/// A collection of [`ArrayReduceRule`]s registered for a specific encoding.
///
/// During optimization, the optimizer calls [`evaluate`](Self::evaluate) which tries each rule
/// in order. The first rule that returns `Some` wins.
pub struct ReduceRuleSet<V: VTable> {
rules: &'static [&'static dyn ArrayReduceRule<V>],
}
impl<V: VTable> ReduceRuleSet<V> {
/// Create a new reduction rule set with the given rules.
pub const fn new(rules: &'static [&'static dyn ArrayReduceRule<V>]) -> Self {
Self { rules }
}
/// Evaluate the reduction rules on the given array.
pub fn evaluate(&self, array: ArrayView<'_, V>) -> VortexResult<Option<ArrayRef>> {
for rule in self.rules.iter() {
if let Some(reduced) = rule.reduce(array)? {
trace_op!(record_reduce_applied(array.array(), *rule, &reduced));
return Ok(Some(reduced));
}
trace_op!(record_reduce_declined(array.array(), *rule));
}
Ok(None)
}
}
/// A set of parent reduction rules for a specific child array encoding.
pub struct ParentRuleSet<V: VTable> {
rules: &'static [&'static dyn DynArrayParentReduceRule<V>],
}
impl<V: VTable> ParentRuleSet<V> {
/// Create a new parent rule set with the given rules.
///
/// Use [`ParentRuleSet::lift`] to lift static rules into dynamic trait objects.
pub const fn new(rules: &'static [&'static dyn DynArrayParentReduceRule<V>]) -> Self {
Self { rules }
}
/// Lift the given rule into a dynamic trait object.
pub const fn lift<R: ArrayParentReduceRule<V>>(
rule: &'static R,
) -> &'static dyn DynArrayParentReduceRule<V> {
// Assert that self is zero-sized
const {
assert!(
!(size_of::<R>() != 0),
"Rule must be zero-sized to be lifted"
);
}
unsafe { &*(rule as *const R as *const ParentReduceRuleAdapter<V, R>) }
}
/// Evaluate the parent reduction rules on the given child and parent arrays.
pub fn evaluate(
&self,
child: ArrayView<'_, V>,
parent: &ArrayRef,
child_idx: usize,
) -> VortexResult<Option<ArrayRef>> {
for rule in self.rules.iter() {
if !rule.matches(parent) {
trace_op!(record_static_parent_reduce_no_match(
parent,
child.array(),
child_idx,
*rule,
));
continue;
}
if let Some(reduced) = rule.reduce_parent(child, parent, child_idx)? {
// Debug assertions because these checks are already run elsewhere.
#[cfg(debug_assertions)]
{
vortex_error::vortex_ensure!(
reduced.len() == parent.len(),
"Reduced array length mismatch from {:?}\nFrom:\n{}\nTo:\n{}",
rule,
parent.encoding_id(),
reduced.encoding_id()
);
vortex_error::vortex_ensure!(
reduced.dtype() == parent.dtype(),
"Reduced array dtype mismatch from {:?}\nFrom:\n{}\nTo:\n{}",
rule,
parent.encoding_id(),
reduced.encoding_id()
);
}
trace_op!(record_static_parent_reduce_applied(
parent,
child.array(),
child_idx,
*rule,
&reduced,
));
return Ok(Some(reduced));
}
trace_op!(record_static_parent_reduce_declined(
parent,
child.array(),
child_idx,
*rule,
));
}
Ok(None)
}
}