Skip to content

Commit 7c1dcaa

Browse files
sddsdd
committed
feat: add InclusiveProjection
1 parent 2de29c2 commit 7c1dcaa

File tree

3 files changed

+506
-35
lines changed

3 files changed

+506
-35
lines changed

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

Lines changed: 131 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,24 @@ 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 => {
207+
visitor.starts_with(reference, literal, predicate)
208+
}
209+
PredicateOperator::NotStartsWith => {
210+
visitor.not_starts_with(reference, literal, predicate)
145211
}
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),
150212
op => {
151213
panic!("Unexpected op for binary predicate: {}", &op)
152214
}
@@ -156,8 +218,8 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
156218
let reference = expr.term();
157219
let literals = expr.literals();
158220
match expr.op() {
159-
PredicateOperator::In => visitor.r#in(reference, literals),
160-
PredicateOperator::NotIn => visitor.not_in(reference, literals),
221+
PredicateOperator::In => visitor.r#in(reference, literals, predicate),
222+
PredicateOperator::NotIn => visitor.not_in(reference, literals, predicate),
161223
op => {
162224
panic!("Unexpected op for set predicate: {}", &op)
163225
}
@@ -170,8 +232,8 @@ pub(crate) fn visit<V: BoundPredicateVisitor>(
170232
mod tests {
171233
use crate::expr::visitors::bound_predicate_visitor::{visit, BoundPredicateVisitor};
172234
use crate::expr::{
173-
BinaryExpression, Bind, BoundReference, Predicate, PredicateOperator, Reference,
174-
SetExpression, UnaryExpression,
235+
BinaryExpression, Bind, BoundPredicate, BoundReference, Predicate, PredicateOperator,
236+
Reference, SetExpression, UnaryExpression,
175237
};
176238
use crate::spec::{Datum, NestedField, PrimitiveType, Schema, SchemaRef, Type};
177239
use fnv::FnvHashSet;
@@ -202,26 +264,43 @@ mod tests {
202264
Ok(!inner)
203265
}
204266

205-
fn is_null(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
267+
fn is_null(
268+
&mut self,
269+
_reference: &BoundReference,
270+
_predicate: &BoundPredicate,
271+
) -> crate::Result<bool> {
206272
Ok(true)
207273
}
208274

209-
fn not_null(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
275+
fn not_null(
276+
&mut self,
277+
_reference: &BoundReference,
278+
_predicate: &BoundPredicate,
279+
) -> crate::Result<bool> {
210280
Ok(false)
211281
}
212282

213-
fn is_nan(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
283+
fn is_nan(
284+
&mut self,
285+
_reference: &BoundReference,
286+
_predicate: &BoundPredicate,
287+
) -> crate::Result<bool> {
214288
Ok(true)
215289
}
216290

217-
fn not_nan(&mut self, _reference: &BoundReference) -> crate::Result<bool> {
291+
fn not_nan(
292+
&mut self,
293+
_reference: &BoundReference,
294+
_predicate: &BoundPredicate,
295+
) -> crate::Result<bool> {
218296
Ok(false)
219297
}
220298

221299
fn less_than(
222300
&mut self,
223301
_reference: &BoundReference,
224302
_literal: &Datum,
303+
_predicate: &BoundPredicate,
225304
) -> crate::Result<bool> {
226305
Ok(true)
227306
}
@@ -230,6 +309,7 @@ mod tests {
230309
&mut self,
231310
_reference: &BoundReference,
232311
_literal: &Datum,
312+
_predicate: &BoundPredicate,
233313
) -> crate::Result<bool> {
234314
Ok(false)
235315
}
@@ -238,6 +318,7 @@ mod tests {
238318
&mut self,
239319
_reference: &BoundReference,
240320
_literal: &Datum,
321+
_predicate: &BoundPredicate,
241322
) -> crate::Result<bool> {
242323
Ok(true)
243324
}
@@ -246,22 +327,34 @@ mod tests {
246327
&mut self,
247328
_reference: &BoundReference,
248329
_literal: &Datum,
330+
_predicate: &BoundPredicate,
249331
) -> crate::Result<bool> {
250332
Ok(false)
251333
}
252334

253-
fn eq(&mut self, _reference: &BoundReference, _literal: &Datum) -> crate::Result<bool> {
335+
fn eq(
336+
&mut self,
337+
_reference: &BoundReference,
338+
_literal: &Datum,
339+
_predicate: &BoundPredicate,
340+
) -> crate::Result<bool> {
254341
Ok(true)
255342
}
256343

257-
fn not_eq(&mut self, _reference: &BoundReference, _literal: &Datum) -> crate::Result<bool> {
344+
fn not_eq(
345+
&mut self,
346+
_reference: &BoundReference,
347+
_literal: &Datum,
348+
_predicate: &BoundPredicate,
349+
) -> crate::Result<bool> {
258350
Ok(false)
259351
}
260352

261353
fn starts_with(
262354
&mut self,
263355
_reference: &BoundReference,
264356
_literal: &Datum,
357+
_predicate: &BoundPredicate,
265358
) -> crate::Result<bool> {
266359
Ok(true)
267360
}
@@ -270,6 +363,7 @@ mod tests {
270363
&mut self,
271364
_reference: &BoundReference,
272365
_literal: &Datum,
366+
_predicate: &BoundPredicate,
273367
) -> crate::Result<bool> {
274368
Ok(false)
275369
}
@@ -278,6 +372,7 @@ mod tests {
278372
&mut self,
279373
_reference: &BoundReference,
280374
_literals: &FnvHashSet<Datum>,
375+
_predicate: &BoundPredicate,
281376
) -> crate::Result<bool> {
282377
Ok(true)
283378
}
@@ -286,6 +381,7 @@ mod tests {
286381
&mut self,
287382
_reference: &BoundReference,
288383
_literals: &FnvHashSet<Datum>,
384+
_predicate: &BoundPredicate,
289385
) -> crate::Result<bool> {
290386
Ok(false)
291387
}

0 commit comments

Comments
 (0)