Skip to content

Commit 442efe2

Browse files
nicklannicklan
authored
feat: Switch scan.execute to return pre-filtered data (delta-io#1429)
## What changes are proposed in this pull request? Switch scan.execute to pre-filter the data. Also adds the needed apis to `EngineData` to support this. I'm encouraged that this leads to a net reduction in code. NB: This is a pretty major public facing change to our APIS (see how many examples/tests had to change). ### This PR affects the following public APIs `scan.execute` ## How was this change tested? Existing tests
1 parent efddf5d commit 442efe2

16 files changed

Lines changed: 131 additions & 340 deletions

File tree

‎acceptance/src/data.rs‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
use std::{path::Path, sync::Arc};
22

33
use delta_kernel::arrow::array::{Array, RecordBatch};
4-
use delta_kernel::arrow::compute::{
5-
concat_batches, filter_record_batch, lexsort_to_indices, take, SortColumn,
6-
};
4+
use delta_kernel::arrow::compute::{concat_batches, lexsort_to_indices, take, SortColumn};
75
use delta_kernel::arrow::datatypes::{DataType, Schema};
86

97
use delta_kernel::parquet::arrow::async_reader::{
@@ -120,23 +118,16 @@ pub async fn assert_scan_metadata(
120118
let mut schema = None;
121119
let batches: Vec<RecordBatch> = scan
122120
.execute(engine)?
123-
.map(|scan_result| -> DeltaResult<_> {
124-
let scan_result = scan_result?;
125-
let mask = scan_result.full_mask();
126-
let data = scan_result.raw_data?;
127-
let record_batch: RecordBatch = data
121+
.map(|data| -> DeltaResult<_> {
122+
let record_batch: RecordBatch = data?
128123
.into_any()
129124
.downcast::<ArrowEngineData>()
130125
.unwrap()
131126
.into();
132127
if schema.is_none() {
133128
schema = Some(record_batch.schema());
134129
}
135-
if let Some(mask) = mask {
136-
Ok(filter_record_batch(&record_batch, &mask.into())?)
137-
} else {
138-
Ok(record_batch)
139-
}
130+
Ok(record_batch)
140131
})
141132
.try_collect()?;
142133
let all_data = concat_batches(&schema.unwrap(), batches.iter()).map_err(Error::from)?;

‎ffi/src/table_changes.rs‎

Lines changed: 7 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,11 @@ use std::sync::Arc;
44
use std::sync::Mutex;
55

66
use delta_kernel::arrow::array::{Array, ArrayData, RecordBatch, StructArray};
7-
use delta_kernel::arrow::compute::filter_record_batch;
87
use delta_kernel::arrow::ffi::to_ffi;
98
use delta_kernel::engine::arrow_data::ArrowEngineData;
10-
use delta_kernel::scan::ScanResult;
119
use delta_kernel::table_changes::scan::TableChangesScan;
1210
use delta_kernel::table_changes::TableChanges;
11+
use delta_kernel::EngineData;
1312
use delta_kernel::Error;
1413
use delta_kernel::{DeltaResult, Version};
1514
use delta_kernel_ffi_macros::handle_descriptor;
@@ -238,8 +237,10 @@ pub unsafe extern "C" fn table_changes_scan_physical_schema(
238237
table_changes_scan.physical_schema().clone().into()
239238
}
240239

240+
type TableChangesData = Mutex<Box<dyn Iterator<Item = DeltaResult<Box<dyn EngineData>>> + Send>>;
241+
241242
pub struct ScanTableChangesIterator {
242-
data: Mutex<Box<dyn Iterator<Item = DeltaResult<ScanResult>> + Send>>,
243+
data: TableChangesData,
243244
engine: Arc<dyn ExternEngine>,
244245
}
245246

@@ -314,22 +315,16 @@ fn scan_table_changes_next_impl(data: &ScanTableChangesIterator) -> DeltaResult<
314315
.lock()
315316
.map_err(|_| Error::generic("poisoned scan table changes iterator mutex"))?;
316317

317-
let Some(scan_result) = data.next().transpose()? else {
318+
let Some(data) = data.next().transpose()? else {
318319
return Ok(ArrowFFIData::empty());
319320
};
320321

321-
let mask = scan_result.full_mask();
322-
let data = scan_result.raw_data?;
323-
let mut record_batch: RecordBatch = data
322+
let record_batch: RecordBatch = data
324323
.into_any()
325324
.downcast::<ArrowEngineData>()
326325
.map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
327326
.into();
328327

329-
if let Some(mask) = mask {
330-
record_batch = filter_record_batch(&record_batch, &mask.into())?;
331-
}
332-
333328
let batch_struct_array: StructArray = record_batch.into();
334329
let array_data: ArrayData = batch_struct_array.into_data();
335330
let (out_array, out_schema) = to_ffi(&array_data)?;
@@ -485,17 +480,7 @@ mod tests {
485480
) -> DeltaResult<Vec<RecordBatch>> {
486481
let scan_results = scan.execute(engine)?;
487482
scan_results
488-
.map(|scan_result| -> DeltaResult<_> {
489-
let scan_result = scan_result?;
490-
let mask = scan_result.full_mask();
491-
let data = scan_result.raw_data?;
492-
let record_batch = to_arrow(data)?;
493-
if let Some(mask) = mask {
494-
Ok(filter_record_batch(&record_batch, &mask.into())?)
495-
} else {
496-
Ok(record_batch)
497-
}
498-
})
483+
.map(|data| -> DeltaResult<_> { to_arrow(data?) })
499484
.try_collect()
500485
}
501486

‎kernel/examples/read-table-changes/src/main.rs‎

Lines changed: 4 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ use std::sync::Arc;
33
use clap::Parser;
44
use common::{LocationArgs, ParseWithExamples};
55
use delta_kernel::arrow::array::RecordBatch;
6-
use delta_kernel::arrow::{compute::filter_record_batch, util::pretty::print_batches};
6+
use delta_kernel::arrow::util::pretty::print_batches;
77
use delta_kernel::engine::arrow_data::ArrowEngineData;
88
use delta_kernel::table_changes::TableChanges;
99
use delta_kernel::DeltaResult;
@@ -38,20 +38,13 @@ fn main() -> DeltaResult<()> {
3838
let table_changes_scan = table_changes.into_scan_builder().build()?;
3939
let batches: Vec<RecordBatch> = table_changes_scan
4040
.execute(Arc::new(engine))?
41-
.map(|scan_result| -> DeltaResult<_> {
42-
let scan_result = scan_result?;
43-
let mask = scan_result.full_mask();
44-
let data = scan_result.raw_data?;
45-
let record_batch: RecordBatch = data
41+
.map(|data| -> DeltaResult<_> {
42+
let record_batch: RecordBatch = data?
4643
.into_any()
4744
.downcast::<ArrowEngineData>()
4845
.map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
4946
.into();
50-
if let Some(mask) = mask {
51-
Ok(filter_record_batch(&record_batch, &mask.into())?)
52-
} else {
53-
Ok(record_batch)
54-
}
47+
Ok(record_batch)
5548
})
5649
.try_collect()?;
5750
print_batches(&batches)?;

‎kernel/examples/read-table-single-threaded/src/main.rs‎

Lines changed: 3 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
use std::process::ExitCode;
22
use std::sync::Arc;
33

4-
use arrow::compute::filter_record_batch;
54
use arrow::record_batch::RecordBatch;
65
use arrow::util::pretty::print_batches;
76
use common::{LocationArgs, ParseWithExamples, ScanArgs};
@@ -50,21 +49,14 @@ fn try_main() -> DeltaResult<()> {
5049
let mut rows_so_far = 0;
5150
let batches: Vec<RecordBatch> = scan
5251
.execute(Arc::new(engine))?
53-
.map(|scan_result| -> DeltaResult<_> {
52+
.map(|data| -> DeltaResult<_> {
5453
// extract the batches and filter them if they have deletion vectors
55-
let scan_result = scan_result?;
56-
let mask = scan_result.full_mask();
57-
let data = scan_result.raw_data?;
58-
let record_batch: RecordBatch = data
54+
let record_batch: RecordBatch = data?
5955
.into_any()
6056
.downcast::<ArrowEngineData>()
6157
.map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
6258
.into();
63-
if let Some(mask) = mask {
64-
Ok(filter_record_batch(&record_batch, &mask.into())?)
65-
} else {
66-
Ok(record_batch)
67-
}
59+
Ok(record_batch)
6860
})
6961
.scan(&mut rows_so_far, |rows_so_far, record_batch| {
7062
// handle truncation if we've specified a limit

‎kernel/examples/write-table/src/main.rs‎

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -320,24 +320,13 @@ async fn read_and_display_data(
320320

321321
let batches: Vec<RecordBatch> = scan
322322
.execute(Arc::new(engine))?
323-
.map(|scan_result| -> DeltaResult<_> {
324-
let scan_result = scan_result?;
325-
let mask = scan_result.full_mask();
326-
let data = scan_result.raw_data?;
327-
let record_batch: RecordBatch = data
323+
.map(|data| -> DeltaResult<_> {
324+
let record_batch: RecordBatch = data?
328325
.into_any()
329326
.downcast::<ArrowEngineData>()
330327
.map_err(|_| Error::EngineDataType("ArrowEngineData".to_string()))?
331328
.into();
332-
333-
if let Some(mask) = mask {
334-
Ok(arrow::compute::filter_record_batch(
335-
&record_batch,
336-
&mask.into(),
337-
)?)
338-
} else {
339-
Ok(record_batch)
340-
}
329+
Ok(record_batch)
341330
})
342331
.try_collect()?;
343332

‎kernel/src/engine/arrow_data.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ use crate::arrow::array::types::{Int32Type, Int64Type};
99
use crate::arrow::array::{
1010
Array, ArrayRef, GenericListArray, MapArray, OffsetSizeTrait, RecordBatch, StructArray,
1111
};
12+
use crate::arrow::compute::filter_record_batch;
1213
use crate::arrow::datatypes::{
1314
DataType as ArrowDataType, Field as ArrowField, FieldRef, Schema as ArrowSchema,
1415
};
@@ -237,6 +238,15 @@ impl EngineData for ArrowEngineData {
237238
let data = RecordBatch::try_new(combined_schema, combined_columns)?;
238239
Ok(Box::new(ArrowEngineData { data }))
239240
}
241+
242+
fn apply_selection_vector(
243+
self: Box<Self>,
244+
mut selection_vector: Vec<bool>,
245+
) -> DeltaResult<Box<dyn EngineData>> {
246+
selection_vector.resize(self.len(), true);
247+
let filtered = filter_record_batch(&self.data, &selection_vector.into())?;
248+
Ok(Box::new(Self::new(filtered)))
249+
}
240250
}
241251

242252
impl ArrowEngineData {

0 commit comments

Comments
 (0)