Skip to content

Commit bd172ae

Browse files
ArnavBalyanArnavBalyan
committed
Add in predicate
1 parent a513a96 commit bd172ae

3 files changed

Lines changed: 193 additions & 0 deletions

File tree

‎core/src/main/java/dev/hardwood/internal/reader/RowGroupFilterEvaluator.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
import dev.hardwood.reader.FilterPredicate.FloatColumnPredicate;
2121
import dev.hardwood.reader.FilterPredicate.IntColumnPredicate;
2222
import dev.hardwood.reader.FilterPredicate.LongColumnPredicate;
23+
import dev.hardwood.reader.FilterPredicate.BinaryInPredicate;
24+
import dev.hardwood.reader.FilterPredicate.IntInPredicate;
25+
import dev.hardwood.reader.FilterPredicate.LongInPredicate;
2326
import dev.hardwood.reader.FilterPredicate.Not;
2427
import dev.hardwood.reader.FilterPredicate.Or;
2528
import dev.hardwood.schema.FileSchema;
@@ -46,6 +49,9 @@ public static boolean canDropRowGroup(FilterPredicate predicate, RowGroup rowGro
4649
case DoubleColumnPredicate p -> evaluateDouble(p, rowGroup, schema);
4750
case BooleanColumnPredicate p -> evaluateBoolean(p, rowGroup, schema);
4851
case BinaryColumnPredicate p -> evaluateBinary(p, rowGroup, schema);
52+
case IntInPredicate p -> evaluateIntIn(p, rowGroup, schema);
53+
case LongInPredicate p -> evaluateLongIn(p, rowGroup, schema);
54+
case BinaryInPredicate p -> evaluateBinaryIn(p, rowGroup, schema);
4955
case And a -> {
5056
for (FilterPredicate f : a.filters()) {
5157
if (canDropRowGroup(f, rowGroup, schema)) {
@@ -195,6 +201,52 @@ private static boolean evaluateBinary(BinaryColumnPredicate p, RowGroup rowGroup
195201
StatisticsDecoder.compareBinary(min, max));
196202
}
197203

204+
private static boolean evaluateIntIn(IntInPredicate p, RowGroup rowGroup, FileSchema schema) {
205+
Statistics stats = findStatistics(p.column(), rowGroup, schema);
206+
if (stats == null || stats.minValue() == null || stats.maxValue() == null) {
207+
return false;
208+
}
209+
int min = StatisticsDecoder.decodeInt(stats.minValue());
210+
int max = StatisticsDecoder.decodeInt(stats.maxValue());
211+
for (int value : p.values()) {
212+
if (value >= min && value <= max) {
213+
return false;
214+
}
215+
}
216+
return true;
217+
}
218+
219+
private static boolean evaluateLongIn(LongInPredicate p, RowGroup rowGroup, FileSchema schema) {
220+
Statistics stats = findStatistics(p.column(), rowGroup, schema);
221+
if (stats == null || stats.minValue() == null || stats.maxValue() == null) {
222+
return false;
223+
}
224+
long min = StatisticsDecoder.decodeLong(stats.minValue());
225+
long max = StatisticsDecoder.decodeLong(stats.maxValue());
226+
for (long value : p.values()) {
227+
if (value >= min && value <= max) {
228+
return false;
229+
}
230+
}
231+
return true;
232+
}
233+
234+
private static boolean evaluateBinaryIn(BinaryInPredicate p, RowGroup rowGroup, FileSchema schema) {
235+
Statistics stats = findStatistics(p.column(), rowGroup, schema);
236+
if (stats == null || stats.minValue() == null || stats.maxValue() == null) {
237+
return false;
238+
}
239+
byte[] min = stats.minValue();
240+
byte[] max = stats.maxValue();
241+
for (byte[] value : p.values()) {
242+
if (StatisticsDecoder.compareBinary(value, min) >= 0
243+
&& StatisticsDecoder.compareBinary(value, max) <= 0) {
244+
return false;
245+
}
246+
}
247+
return true;
248+
}
249+
198250
// ==================== Generic comparison logic ====================
199251

200252
/// Determines if a row group can be dropped given integer-comparable min/max statistics.

‎core/src/main/java/dev/hardwood/reader/FilterPredicate.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ public sealed interface FilterPredicate
4040
FilterPredicate.DoubleColumnPredicate,
4141
FilterPredicate.BooleanColumnPredicate,
4242
FilterPredicate.BinaryColumnPredicate,
43+
FilterPredicate.IntInPredicate,
44+
FilterPredicate.LongInPredicate,
45+
FilterPredicate.BinaryInPredicate,
4346
FilterPredicate.And,
4447
FilterPredicate.Or,
4548
FilterPredicate.Not {
@@ -190,6 +193,22 @@ static FilterPredicate gtEq(String column, String value) {
190193
return new BinaryColumnPredicate(column, Operator.GT_EQ, value.getBytes(StandardCharsets.UTF_8));
191194
}
192195

196+
static FilterPredicate in(String column, int... values) {
197+
return new IntInPredicate(column, values);
198+
}
199+
200+
static FilterPredicate in(String column, long... values) {
201+
return new LongInPredicate(column, values);
202+
}
203+
204+
static FilterPredicate inStrings(String column, String... values) {
205+
byte[][] encoded = new byte[values.length][];
206+
for (int i = 0; i < values.length; i++) {
207+
encoded[i] = values[i].getBytes(StandardCharsets.UTF_8);
208+
}
209+
return new BinaryInPredicate(column, encoded);
210+
}
211+
193212
// ==================== Logical Combinators ====================
194213

195214
static FilterPredicate and(FilterPredicate left, FilterPredicate right) {
@@ -247,6 +266,51 @@ public int hashCode() {
247266
}
248267
}
249268

269+
record IntInPredicate(String column, int[] values) implements FilterPredicate {
270+
271+
@Override
272+
public boolean equals(Object o) {
273+
if (this == o) return true;
274+
if (!(o instanceof IntInPredicate that)) return false;
275+
return column.equals(that.column) && Arrays.equals(values, that.values);
276+
}
277+
278+
@Override
279+
public int hashCode() {
280+
return 31 * column.hashCode() + Arrays.hashCode(values);
281+
}
282+
}
283+
284+
record LongInPredicate(String column, long[] values) implements FilterPredicate {
285+
286+
@Override
287+
public boolean equals(Object o) {
288+
if (this == o) return true;
289+
if (!(o instanceof LongInPredicate that)) return false;
290+
return column.equals(that.column) && Arrays.equals(values, that.values);
291+
}
292+
293+
@Override
294+
public int hashCode() {
295+
return 31 * column.hashCode() + Arrays.hashCode(values);
296+
}
297+
}
298+
299+
record BinaryInPredicate(String column, byte[][] values) implements FilterPredicate {
300+
301+
@Override
302+
public boolean equals(Object o) {
303+
if (this == o) return true;
304+
if (!(o instanceof BinaryInPredicate that)) return false;
305+
return column.equals(that.column) && Arrays.deepEquals(values, that.values);
306+
}
307+
308+
@Override
309+
public int hashCode() {
310+
return 31 * column.hashCode() + Arrays.deepHashCode(values);
311+
}
312+
}
313+
250314
// ==================== Logical Combinator Records ====================
251315

252316
record And(List<FilterPredicate> filters) implements FilterPredicate {

‎core/src/test/java/dev/hardwood/FilterPredicateTest.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -370,6 +370,83 @@ void testUnknownColumnNeverDrop() {
370370
FilterPredicate.eq("nonexistent", 42), rg, schema)).isFalse();
371371
}
372372

373+
@Test
374+
void testIntInPredicateCreation() {
375+
FilterPredicate p = FilterPredicate.in("id", 1, 5, 10);
376+
assertThat(p).isInstanceOf(FilterPredicate.IntInPredicate.class);
377+
FilterPredicate.IntInPredicate ip = (FilterPredicate.IntInPredicate) p;
378+
assertThat(ip.column()).isEqualTo("id");
379+
assertThat(ip.values()).containsExactly(1, 5, 10);
380+
}
381+
382+
@Test
383+
void testLongInPredicateCreation() {
384+
FilterPredicate p = FilterPredicate.in("ts", 100L, 200L);
385+
assertThat(p).isInstanceOf(FilterPredicate.LongInPredicate.class);
386+
}
387+
388+
@Test
389+
void testStringInPredicateCreation() {
390+
FilterPredicate p = FilterPredicate.inStrings("city", "NYC", "LA");
391+
assertThat(p).isInstanceOf(FilterPredicate.BinaryInPredicate.class);
392+
}
393+
394+
@Test
395+
void testCanDropWithIntIn() {
396+
RowGroup rg = createIntRowGroup(10, 20);
397+
FileSchema schema = createIntSchema();
398+
399+
// All values below min -> can drop
400+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
401+
FilterPredicate.in("col", 1, 5, 8), rg, schema)).isTrue();
402+
403+
// All values above max -> can drop
404+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
405+
FilterPredicate.in("col", 25, 30), rg, schema)).isTrue();
406+
407+
// One value in range -> cannot drop
408+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
409+
FilterPredicate.in("col", 5, 15, 25), rg, schema)).isFalse();
410+
411+
// Value on boundary -> cannot drop
412+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
413+
FilterPredicate.in("col", 1, 10), rg, schema)).isFalse();
414+
}
415+
416+
@Test
417+
void testCanDropWithLongIn() {
418+
RowGroup rg = createLongRowGroup(100L, 200L);
419+
FileSchema schema = createLongSchema();
420+
421+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
422+
FilterPredicate.in("col", 50L, 80L), rg, schema)).isTrue();
423+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
424+
FilterPredicate.in("col", 50L, 150L), rg, schema)).isFalse();
425+
}
426+
427+
@Test
428+
void testCanDropWithStringIn() {
429+
RowGroup rg = createBinaryRowGroup("banana".getBytes(), "date".getBytes());
430+
FileSchema schema = createBinarySchema();
431+
432+
// All outside range -> can drop
433+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
434+
FilterPredicate.inStrings("col", "apple", "elderberry"), rg, schema)).isTrue();
435+
436+
// "cherry" in range -> cannot drop
437+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
438+
FilterPredicate.inStrings("col", "apple", "cherry"), rg, schema)).isFalse();
439+
}
440+
441+
@Test
442+
void testCanDropWithInMissingStatistics() {
443+
RowGroup rg = createRowGroupWithoutStatistics();
444+
FileSchema schema = createIntSchema();
445+
446+
assertThat(RowGroupFilterEvaluator.canDropRowGroup(
447+
FilterPredicate.in("col", 1, 2, 3), rg, schema)).isFalse();
448+
}
449+
373450
// ==================== Helpers ====================
374451

375452
private static RowGroup createIntRowGroup(int min, int max) {

0 commit comments

Comments
 (0)