Skip to content

Commit 95f1b22

Browse files
sddsdd
committed
feat: add InclusiveProjection
1 parent 9160bba commit 95f1b22

File tree

3 files changed

+574
-35
lines changed

3 files changed

+574
-35
lines changed

‎crates/iceberg/src/expr/visitors/bound_predicate_visitor.rs‎

Lines changed: 129 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -41,54 +41,108 @@ pub trait BoundPredicateVisitor {
4141
fn not(&mut self, inner: Self::T) -> Result<Self::T>;
4242

4343
/// Called after a predicate with an `IsNull` operator is visited
44-
fn is_null(&mut self, reference: &BoundReference) -> Result<Self::T>;
44+
fn is_null(
45+
&mut self,
46+
reference: &BoundReference,
47+
predicate: &BoundPredicate,
48+
) -> Result<Self::T>;
4549

4650
/// Called after a predicate with a `NotNull` operator is visited
47-
fn not_null(&mut self, reference: &BoundReference) -> Result<Self::T>;
51+
fn not_null(
52+
&mut self,
53+
reference: &BoundReference,
54+
predicate: &BoundPredicate,
55+
) -> Result<Self::T>;
4856

4957
/// Called after a predicate with an `IsNan` operator is visited
50-
fn is_nan(&mut self, reference: &BoundReference) -> Result<Self::T>;
58+
fn is_nan(&mut self, reference: &BoundReference, predicate: &BoundPredicate)
59+
-> Result<Self::T>;
5160

5261
/// Called after a predicate with a `NotNan` operator is visited
53-
fn not_nan(&mut self, reference: &BoundReference) -> Result<Self::T>;
62+
fn not_nan(
63+
&mut self,
64+
reference: &BoundReference,
65+
predicate: &BoundPredicate,
66+
) -> Result<Self::T>;
5467

5568
/// Called after a predicate with a `LessThan` operator is visited
56-
fn less_than(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
69+
fn less_than(
70+
&mut self,
71+
reference: &BoundReference,
72+
literal: &Datum,
73+
predicate: &BoundPredicate,
74+
) -> Result<Self::T>;
5775

5876
/// Called after a predicate with a `LessThanOrEq` operator is visited
59-
fn less_than_or_eq(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
77+
fn less_than_or_eq(
78+
&mut self,
79+
reference: &BoundReference,
80+
literal: &Datum,
81+
predicate: &BoundPredicate,
82+
) -> Result<Self::T>;
6083

6184
/// Called after a predicate with a `GreaterThan` operator is visited
62-
fn greater_than(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
85+
fn greater_than(
86+
&mut self,
87+
reference: &BoundReference,
88+
literal: &Datum,
89+
predicate: &BoundPredicate,
90+
) -> Result<Self::T>;
6391

6492
/// Called after a predicate with a `GreaterThanOrEq` operator is visited
6593
fn greater_than_or_eq(
6694
&mut self,
6795
reference: &BoundReference,
6896
literal: &Datum,
97+
predicate: &BoundPredicate,
6998
) -> Result<Self::T>;
7099

71100
/// Called after a predicate with an `Eq` operator is visited
72-
fn eq(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
101+
fn eq(
102+
&mut self,
103+
reference: &BoundReference,
104+
literal: &Datum,
105+
predicate: &BoundPredicate,
106+
) -> Result<Self::T>;
73107

74108
/// Called after a predicate with a `NotEq` operator is visited
75-
fn not_eq(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
109+
fn not_eq(
110+
&mut self,
111+
reference: &BoundReference,
112+
literal: &Datum,
113+
predicate: &BoundPredicate,
114+
) -> Result<Self::T>;
76115

77116
/// Called after a predicate with a `StartsWith` operator is visited
78-
fn starts_with(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
117+
fn starts_with(
118+
&mut self,
119+
reference: &BoundReference,
120+
literal: &Datum,
121+
predicate: &BoundPredicate,
122+
) -> Result<Self::T>;
79123

80124
/// Called after a predicate with a `NotStartsWith` operator is visited
81-
fn not_starts_with(&mut self, reference: &BoundReference, literal: &Datum) -> Result<Self::T>;
125+
fn not_starts_with(
126+
&mut self,
127+
reference: &BoundReference,
128+
literal: &Datum,
129+
predicate: &BoundPredicate,
130+
) -> Result<Self::T>;
82131

83132
/// Called after a predicate with an `In` operator is visited
84-
fn r#in(&mut self, reference: &BoundReference, literals: &FnvHashSet<Datum>)
85-
-> Result<Self::T>;
133+
fn r#in(
134+
&mut self,
135+
reference: &BoundReference,
136+
literals: &FnvHashSet<Datum>,
137+
predicate: &BoundPredicate,
138+
) -> Result<Self::T>;
86139

87140
/// Called after a predicate with a `NotIn` operator is visited
88141
fn not_in(
89142
&mut self,
90143
reference: &BoundReference,
91144
literals: &FnvHashSet<Datum>,
145+
predicate: &BoundPredicate,
92146
) -> Result<Self::T>;
93147
}
94148

@@ -125,10 +179,10 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
125179
visitor.not(inner_result)
126180
}
127181
BoundPredicate::Unary(expr) => match expr.op() {
128-
PredicateOperator::IsNull => visitor.is_null(expr.term()),
129-
PredicateOperator::NotNull => visitor.not_null(expr.term()),
130-
PredicateOperator::IsNan => visitor.is_nan(expr.term()),
131-
PredicateOperator::NotNan => visitor.not_nan(expr.term()),
182+
PredicateOperator::IsNull => visitor.is_null(expr.term(), predicate),
183+
PredicateOperator::NotNull => visitor.not_null(expr.term(), predicate),
184+
PredicateOperator::IsNan => visitor.is_nan(expr.term(), predicate),
185+
PredicateOperator::NotNan => visitor.not_nan(expr.term(), predicate),
132186
op => {
133187
panic!("Unexpected op for unary predicate: {}", &op)
134188
}
@@ -137,16 +191,22 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
137191
let reference = expr.term();
138192
let literal = expr.literal();
139193
match expr.op() {
140-
PredicateOperator::LessThan => visitor.less_than(reference, literal),
141-
PredicateOperator::LessThanOrEq => visitor.less_than_or_eq(reference, literal),
142-
PredicateOperator::GreaterThan => visitor.greater_than(reference, literal),
194+
PredicateOperator::LessThan => visitor.less_than(reference, literal, predicate),
195+
PredicateOperator::LessThanOrEq => {
196+
visitor.less_than_or_eq(reference, literal, predicate)
197+
}
198+
PredicateOperator::GreaterThan => {
199+
visitor.greater_than(reference, literal, predicate)
200+
}
143201
PredicateOperator::GreaterThanOrEq => {
144-
visitor.greater_than_or_eq(reference, literal)
202+
visitor.greater_than_or_eq(reference, literal, predicate)
203+
}
204+
PredicateOperator::Eq => visitor.eq(reference, literal, predicate),
205+
PredicateOperator::NotEq => visitor.not_eq(reference, literal, predicate),
206+
PredicateOperator::StartsWith => visitor.starts_with(reference, literal, predicate),
207+
PredicateOperator::NotStartsWith => {
208+
visitor.not_starts_with(reference, literal, predicate)
145209
}
146-
PredicateOperator::Eq => visitor.eq(reference, literal),
147-
PredicateOperator::NotEq => visitor.not_eq(reference, literal),
148-
PredicateOperator::StartsWith => visitor.starts_with(reference, literal),
149-
PredicateOperator::NotStartsWith => visitor.not_starts_with(reference, literal),
150210
op => {
151211
panic!("Unexpected op for binary predicate: {}", &op)
152212
}
@@ -156,8 +216,8 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
156216
let reference = expr.term();
157217
let literals = expr.literals();
158218
match expr.op() {
159-
PredicateOperator::In => visitor.r#in(reference, literals),
160-
PredicateOperator::NotIn => visitor.not_in(reference, literals),
219+
PredicateOperator::In => visitor.r#in(reference, literals, predicate),
220+
PredicateOperator::NotIn => visitor.not_in(reference, literals, predicate),
161221
op => {
162222
panic!("Unexpected op for set predicate: {}", &op)
163223
}
@@ -170,8 +230,8 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
170230
mod tests {
171231
use crate::expr::visitors::bound_predicate_visitor::{visit, BoundPredicateVisitor};
172232
use crate::expr::{
173-
BinaryExpression, Bind, BoundReference, Predicate, PredicateOperator, Reference,
174-
SetExpression, UnaryExpression,
233+
BinaryExpression, Bind, BoundPredicate, BoundReference, Predicate, PredicateOperator,
234+
Reference, SetExpression, UnaryExpression,
175235
};
176236
use crate::spec::{Datum, NestedField, PrimitiveType, Schema, SchemaRef, Type};
177237
use fnv::FnvHashSet;
@@ -202,26 +262,43 @@ mod tests {
202262
Ok(!inner)
203263
}
204264

205-
fn is_null(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
265+
fn is_null(
266+
&mut self,
267+
_reference: &BoundReference,
268+
_predicate: &BoundPredicate,
269+
) -> crate::Result<bool> {
206270
Ok(true)
207271
}
208272

209-
fn not_null(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
273+
fn not_null(
274+
&mut self,
275+
_reference: &BoundReference,
276+
_predicate: &BoundPredicate,
277+
) -> crate::Result<bool> {
210278
Ok(false)
211279
}
212280

213-
fn is_nan(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
281+
fn is_nan(
282+
&mut self,
283+
_reference: &BoundReference,
284+
_predicate: &BoundPredicate,
285+
) -> crate::Result<bool> {
214286
Ok(true)
215287
}
216288

217-
fn not_nan(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
289+
fn not_nan(
290+
&mut self,
291+
_reference: &BoundReference,
292+
_predicate: &BoundPredicate,
293+
) -> crate::Result<bool> {
218294
Ok(false)
219295
}
220296

221297
fn less_than(
222298
&mut self,
223299
_reference: &BoundReference,
224300
_literal: &Datum,
301+
_predicate: &BoundPredicate,
225302
) -> crate::Result<bool> {
226303
Ok(true)
227304
}
@@ -230,6 +307,7 @@ mod tests {
230307
&mut self,
231308
_reference: &BoundReference,
232309
_literal: &Datum,
310+
_predicate: &BoundPredicate,
233311
) -> crate::Result<bool> {
234312
Ok(false)
235313
}
@@ -238,6 +316,7 @@ mod tests {
238316
&mut self,
239317
_reference: &BoundReference,
240318
_literal: &Datum,
319+
_predicate: &BoundPredicate,
241320
) -> crate::Result<bool> {
242321
Ok(true)
243322
}
@@ -246,22 +325,34 @@ mod tests {
246325
&mut self,
247326
_reference: &BoundReference,
248327
_literal: &Datum,
328+
_predicate: &BoundPredicate,
249329
) -> crate::Result<bool> {
250330
Ok(false)
251331
}
252332

253-
fn eq(&mut self, _reference: &BoundReference, _literal: &Datum) -> crate::Result<bool> {
333+
fn eq(
334+
&mut self,
335+
_reference: &BoundReference,
336+
_literal: &Datum,
337+
_predicate: &BoundPredicate,
338+
) -> crate::Result<bool> {
254339
Ok(true)
255340
}
256341

257-
fn not_eq(&mut self, _reference: &BoundReference, _literal: &Datum) -> crate::Result<bool> {
342+
fn not_eq(
343+
&mut self,
344+
_reference: &BoundReference,
345+
_literal: &Datum,
346+
_predicate: &BoundPredicate,
347+
) -> crate::Result<bool> {
258348
Ok(false)
259349
}
260350

261351
fn starts_with(
262352
&mut self,
263353
_reference: &BoundReference,
264354
_literal: &Datum,
355+
_predicate: &BoundPredicate,
265356
) -> crate::Result<bool> {
266357
Ok(true)
267358
}
@@ -270,6 +361,7 @@ mod tests {
270361
&mut self,
271362
_reference: &BoundReference,
272363
_literal: &Datum,
364+
_predicate: &BoundPredicate,
273365
) -> crate::Result<bool> {
274366
Ok(false)
275367
}
@@ -278,6 +370,7 @@ mod tests {
278370
&mut self,
279371
_reference: &BoundReference,
280372
_literals: &FnvHashSet<Datum>,
373+
_predicate: &BoundPredicate,
281374
) -> crate::Result<bool> {
282375
Ok(true)
283376
}
@@ -286,6 +379,7 @@ mod tests {
286379
&mut self,
287380
_reference: &BoundReference,
288381
_literals: &FnvHashSet<Datum>,
382+
_predicate: &BoundPredicate,
289383
) -> crate::Result<bool> {
290384
Ok(false)
291385
}

0 commit comments

Comments
 (0)