Skip to content

Commit 0b06463

Browse files
shardulm94shardulm94
committed
Row level filtering: Allow table scans to pass a row level filter for ORC files
1 parent 9298ca2 commit 0b06463

File tree

11 files changed

+572
-28
lines changed

11 files changed

+572
-28
lines changed

‎core/src/main/java/org/apache/iceberg/BaseFileScanTask.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -169,7 +169,7 @@ public FileScanTask next() {
169169
}
170170
}
171171

172-
private static final class SplitScanTask implements FileScanTask {
172+
public static final class SplitScanTask implements FileScanTask {
173173
private final long len;
174174
private final long offset;
175175
private final FileScanTask fileScanTask;
@@ -209,5 +209,9 @@ public Expression residual() {
209209
public Iterable<FileScanTask> split(long splitSize) {
210210
throw new UnsupportedOperationException("Cannot split a task which is already split");
211211
}
212+
213+
public FileScanTask underlyingFileScanTask() {
214+
return fileScanTask;
215+
}
212216
}
213217
}

‎data/src/test/java/org/apache/iceberg/data/orc/TestOrcRowLevelFiltering.java‎

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.iceberg.data.orc;
21+
22+
import java.io.File;
23+
import java.io.IOException;
24+
import java.util.List;
25+
import java.util.stream.Collectors;
26+
import java.util.stream.LongStream;
27+
import org.apache.iceberg.Files;
28+
import org.apache.iceberg.Schema;
29+
import org.apache.iceberg.data.DataTestHelpers;
30+
import org.apache.iceberg.data.GenericRecord;
31+
import org.apache.iceberg.data.Record;
32+
import org.apache.iceberg.io.CloseableIterable;
33+
import org.apache.iceberg.io.FileAppender;
34+
import org.apache.iceberg.orc.ORC;
35+
import org.apache.iceberg.orc.OrcRowFilter;
36+
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
37+
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
38+
import org.apache.iceberg.types.Types;
39+
import org.junit.Assert;
40+
import org.junit.Rule;
41+
import org.junit.Test;
42+
import org.junit.rules.TemporaryFolder;
43+
44+
import static org.apache.iceberg.types.Types.NestedField.optional;
45+
import static org.apache.iceberg.types.Types.NestedField.required;
46+
47+
public class TestOrcRowLevelFiltering {
48+
49+
@Rule
50+
public TemporaryFolder temp = new TemporaryFolder();
51+
52+
private static final Schema SCHEMA = new Schema(
53+
required(100, "id", Types.LongType.get()),
54+
required(101, "data1", Types.StringType.get()),
55+
required(102, "data2", Types.StringType.get())
56+
);
57+
58+
private static final List<Record> RECORDS = LongStream.range(0, 100).mapToObj(i -> {
59+
Record record = GenericRecord.create(SCHEMA);
60+
record.set(0, i);
61+
record.set(1, "data1:" + i);
62+
record.set(2, "data2:" + i);
63+
return record;
64+
}).collect(Collectors.toList());
65+
66+
@Test
67+
public void testReadOrcWithRowFilterNoProjection() throws IOException {
68+
testReadOrcWithRowFilter(SCHEMA, rowFilterId(), RECORDS.subList(75, 100));
69+
}
70+
71+
@Test
72+
public void testReadOrcWithRowFilterProjection() throws IOException {
73+
Schema projectedSchema = new Schema(
74+
required(101, "data1", Types.StringType.get())
75+
);
76+
77+
List<Record> expected = RECORDS.subList(75, 100).stream().map(r -> {
78+
Record record = GenericRecord.create(projectedSchema);
79+
record.set(0, r.get(1));
80+
return record;
81+
}).collect(Collectors.toList());
82+
83+
testReadOrcWithRowFilter(projectedSchema, rowFilterId(), expected);
84+
}
85+
86+
@Test
87+
public void testReadOrcWithRowFilterPartialFilterColumns() throws IOException {
88+
Schema projectedSchema = new Schema(
89+
required(101, "data1", Types.StringType.get()),
90+
required(102, "data2", Types.StringType.get())
91+
);
92+
93+
List<Record> expected = RECORDS.subList(25, 75).stream().map(r -> {
94+
Record record = GenericRecord.create(projectedSchema);
95+
record.set(0, r.get(1));
96+
record.set(1, r.get(2));
97+
return record;
98+
}).collect(Collectors.toList());
99+
100+
testReadOrcWithRowFilter(projectedSchema, rowFilterIdAndData1(), expected);
101+
}
102+
103+
@Test
104+
public void testReadOrcWithRowFilterNonExistentColumn() throws IOException {
105+
testReadOrcWithRowFilter(SCHEMA, rowFilterData3(), ImmutableList.of());
106+
}
107+
108+
private void testReadOrcWithRowFilter(Schema schema, OrcRowFilter rowFilter, List<Record> expected)
109+
throws IOException {
110+
File testFile = temp.newFile();
111+
Assert.assertTrue("Delete should succeed", testFile.delete());
112+
try (FileAppender<Record> writer = ORC.write(Files.localOutput(testFile))
113+
.schema(SCHEMA)
114+
.createWriterFunc(GenericOrcWriter::buildWriter)
115+
.build()) {
116+
for (Record rec : RECORDS) {
117+
writer.add(rec);
118+
}
119+
}
120+
121+
List<Record> rows;
122+
try (CloseableIterable<Record> reader = ORC.read(Files.localInput(testFile))
123+
.project(schema)
124+
.createReaderFunc(fileSchema -> GenericOrcReader.buildReader(schema, fileSchema))
125+
.rowFilter(rowFilter)
126+
.build()) {
127+
rows = Lists.newArrayList(reader);
128+
}
129+
130+
for (int i = 0; i < expected.size(); i += 1) {
131+
DataTestHelpers.assertEquals(schema.asStruct(), expected.get(i), rows.get(i));
132+
}
133+
}
134+
135+
private OrcRowFilter rowFilterId() {
136+
return new OrcRowFilter() {
137+
@Override
138+
public Schema requiredSchema() {
139+
return new Schema(
140+
required(100, "id", Types.LongType.get())
141+
);
142+
}
143+
144+
@Override
145+
public boolean shouldKeep(Object[] values) {
146+
return (Long) values[0] >= 75;
147+
}
148+
};
149+
}
150+
151+
private OrcRowFilter rowFilterIdAndData1() {
152+
return new OrcRowFilter() {
153+
@Override
154+
public Schema requiredSchema() {
155+
return new Schema(
156+
SCHEMA.findField("id"),
157+
SCHEMA.findField("data1")
158+
);
159+
}
160+
161+
@Override
162+
public boolean shouldKeep(Object[] values) {
163+
return (Long) values[0] >= 25 && ((String) values[1]).compareTo("data1:75") < 0;
164+
}
165+
};
166+
}
167+
168+
private OrcRowFilter rowFilterData3() {
169+
return new OrcRowFilter() {
170+
@Override
171+
public Schema requiredSchema() {
172+
return new Schema(
173+
optional(104, "data3", Types.LongType.get())
174+
);
175+
}
176+
177+
@Override
178+
public boolean shouldKeep(Object[] values) {
179+
return values[0] != null && (Long) values[0] >= 25;
180+
}
181+
};
182+
}
183+
}

‎orc/src/main/java/org/apache/iceberg/orc/ORC.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@ public static class ReadBuilder {
126126
private Long length = null;
127127
private Expression filter = null;
128128
private boolean caseSensitive = true;
129+
private OrcRowFilter rowFilter = null;
129130

130131
private Function<TypeDescription, OrcRowReader<?>> readerFunc;
131132
private Function<TypeDescription, OrcBatchReader<?>> batchedReaderFunc;
@@ -194,10 +195,15 @@ public ReadBuilder recordsPerBatch(int numRecordsPerBatch) {
194195
return this;
195196
}
196197

198+
public ReadBuilder rowFilter(OrcRowFilter newRowFilter) {
199+
this.rowFilter = newRowFilter;
200+
return this;
201+
}
202+
197203
public <D> CloseableIterable<D> build() {
198204
Preconditions.checkNotNull(schema, "Schema is required");
199205
return new OrcIterable<>(file, conf, schema, start, length, readerFunc, caseSensitive, filter, batchedReaderFunc,
200-
recordsPerBatch);
206+
recordsPerBatch, rowFilter);
201207
}
202208
}
203209

‎orc/src/main/java/org/apache/iceberg/orc/OrcIterable.java‎

Lines changed: 85 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
package org.apache.iceberg.orc;
2121

2222
import java.io.IOException;
23+
import java.util.Set;
2324
import java.util.function.Function;
2425
import org.apache.hadoop.conf.Configuration;
2526
import org.apache.iceberg.Schema;
@@ -31,6 +32,9 @@
3132
import org.apache.iceberg.io.CloseableIterable;
3233
import org.apache.iceberg.io.CloseableIterator;
3334
import org.apache.iceberg.io.InputFile;
35+
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
36+
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
37+
import org.apache.iceberg.types.TypeUtil;
3438
import org.apache.iceberg.util.Pair;
3539
import org.apache.orc.Reader;
3640
import org.apache.orc.TypeDescription;
@@ -51,11 +55,13 @@ class OrcIterable<T> extends CloseableGroup implements CloseableIterable<T> {
5155
private final boolean caseSensitive;
5256
private final Function<TypeDescription, OrcBatchReader<?>> batchReaderFunction;
5357
private final int recordsPerBatch;
58+
private final OrcRowFilter rowFilter;
5459

5560
OrcIterable(InputFile file, Configuration config, Schema schema,
5661
Long start, Long length,
5762
Function<TypeDescription, OrcRowReader<?>> readerFunction, boolean caseSensitive, Expression filter,
58-
Function<TypeDescription, OrcBatchReader<?>> batchReaderFunction, int recordsPerBatch) {
63+
Function<TypeDescription, OrcBatchReader<?>> batchReaderFunction, int recordsPerBatch,
64+
OrcRowFilter rowFilter) {
5965
this.schema = schema;
6066
this.readerFunction = readerFunction;
6167
this.file = file;
@@ -66,6 +72,7 @@ class OrcIterable<T> extends CloseableGroup implements CloseableIterable<T> {
6672
this.filter = (filter == Expressions.alwaysTrue()) ? null : filter;
6773
this.batchReaderFunction = batchReaderFunction;
6874
this.recordsPerBatch = recordsPerBatch;
75+
this.rowFilter = rowFilter;
6976
}
7077

7178
@SuppressWarnings("unchecked")
@@ -81,16 +88,37 @@ public CloseableIterator<T> iterator() {
8188
sarg = ExpressionToSearchArgument.convert(boundFilter, readOrcSchema);
8289
}
8390

84-
VectorizedRowBatchIterator rowBatchIterator = newOrcIterator(file, readOrcSchema, start, length, orcFileReader,
85-
sarg, recordsPerBatch);
86-
if (batchReaderFunction != null) {
87-
OrcBatchReader<T> batchReader = (OrcBatchReader<T>) batchReaderFunction.apply(readOrcSchema);
88-
return CloseableIterator.transform(rowBatchIterator, pair -> {
89-
batchReader.setBatchContext(pair.second());
90-
return batchReader.read(pair.first());
91-
});
91+
if (rowFilter == null) {
92+
VectorizedRowBatchIterator rowBatchIterator = newOrcIterator(file, readOrcSchema, start, length, orcFileReader,
93+
sarg, recordsPerBatch);
94+
if (batchReaderFunction != null) {
95+
OrcBatchReader<T> batchReader = (OrcBatchReader<T>) batchReaderFunction.apply(readOrcSchema);
96+
return CloseableIterator.transform(rowBatchIterator, pair -> {
97+
batchReader.setBatchContext(pair.second());
98+
return batchReader.read(pair.first());
99+
});
100+
} else {
101+
return new OrcRowIterator<>(rowBatchIterator, (OrcRowReader<T>) readerFunction.apply(readOrcSchema),
102+
null, null);
103+
}
92104
} else {
93-
return new OrcRowIterator<>(rowBatchIterator, (OrcRowReader<T>) readerFunction.apply(readOrcSchema));
105+
Preconditions.checkArgument(batchReaderFunction == null,
106+
"Row-level filtering not supported by vectorized reader");
107+
Set<Integer> filterColumnIds = TypeUtil.getProjectedIds(rowFilter.requiredSchema());
108+
Set<Integer> filterColumnIdsNotInReadSchema = Sets.difference(filterColumnIds,
109+
TypeUtil.getProjectedIds(schema));
110+
Schema extraFilterColumns = TypeUtil.select(rowFilter.requiredSchema(), filterColumnIdsNotInReadSchema);
111+
Schema finalReadSchema = TypeUtil.join(schema, extraFilterColumns);
112+
113+
TypeDescription finalReadOrcSchema = ORCSchemaUtil.buildOrcProjection(finalReadSchema,
114+
orcFileReader.getSchema());
115+
TypeDescription rowFilterOrcSchema = ORCSchemaUtil.buildOrcProjection(rowFilter.requiredSchema(),
116+
orcFileReader.getSchema());
117+
RowFilterValueReader filterReader = new RowFilterValueReader(finalReadOrcSchema, rowFilterOrcSchema);
118+
119+
return new OrcRowIterator<>(
120+
newOrcIterator(file, finalReadOrcSchema, start, length, orcFileReader, sarg, recordsPerBatch),
121+
(OrcRowReader<T>) readerFunction.apply(readOrcSchema), rowFilter, filterReader);
94122
}
95123
}
96124

@@ -116,34 +144,67 @@ private static VectorizedRowBatchIterator newOrcIterator(InputFile file,
116144

117145
private static class OrcRowIterator<T> implements CloseableIterator<T> {
118146

119-
private int nextRow;
120-
private VectorizedRowBatch current;
147+
private int currentRow;
148+
private VectorizedRowBatch currentBatch;
149+
private boolean advanced = false;
121150

122151
private final VectorizedRowBatchIterator batchIter;
123152
private final OrcRowReader<T> reader;
153+
private final OrcRowFilter filter;
154+
private final RowFilterValueReader filterReader;
124155

125-
OrcRowIterator(VectorizedRowBatchIterator batchIter, OrcRowReader<T> reader) {
156+
OrcRowIterator(VectorizedRowBatchIterator batchIter, OrcRowReader<T> reader, OrcRowFilter filter,
157+
RowFilterValueReader filterReader) {
126158
this.batchIter = batchIter;
127159
this.reader = reader;
128-
current = null;
129-
nextRow = 0;
160+
this.filter = filter;
161+
this.filterReader = filterReader;
162+
currentBatch = null;
163+
currentRow = 0;
164+
}
165+
166+
private void advance() {
167+
if (!advanced) {
168+
while (true) {
169+
currentRow++;
170+
// if batch has been consumed, move to next batch
171+
if (currentBatch == null || currentRow >= currentBatch.size) {
172+
if (batchIter.hasNext()) {
173+
Pair<VectorizedRowBatch, Long> nextBatch = batchIter.next();
174+
currentBatch = nextBatch.first();
175+
currentRow = 0;
176+
reader.setBatchContext(nextBatch.second());
177+
if (filterReader != null) {
178+
filterReader.setBatchContext(nextBatch.second());
179+
}
180+
} else {
181+
// no more batches left to process
182+
currentBatch = null;
183+
currentRow = -1;
184+
break;
185+
}
186+
}
187+
if (filter == null || filter.shouldKeep(filterReader.read(currentBatch, currentRow))) {
188+
// we have found our row
189+
break;
190+
}
191+
}
192+
advanced = true;
193+
}
130194
}
131195

132196
@Override
133197
public boolean hasNext() {
134-
return (current != null && nextRow < current.size) || batchIter.hasNext();
198+
advance();
199+
return currentBatch != null;
135200
}
136201

137202
@Override
138203
public T next() {
139-
if (current == null || nextRow >= current.size) {
140-
Pair<VectorizedRowBatch, Long> nextBatch = batchIter.next();
141-
current = nextBatch.first();
142-
nextRow = 0;
143-
this.reader.setBatchContext(nextBatch.second());
144-
}
145-
146-
return this.reader.read(current, nextRow++);
204+
advance();
205+
// mark current row as used
206+
advanced = false;
207+
return this.reader.read(currentBatch, currentRow);
147208
}
148209

149210
@Override

0 commit comments

Comments
 (0)