Make deletion vector manipulation less error-prone

acceptance/src/data.rs

@@ -130,6 +130,7 @@ pub async fn assert_scan_data(engine: Arc<dyn Engine>, test_case: &TestCaseInfo)
         .execute(engine)?
         .map(|scan_result| -> DeltaResult<_> {
             let scan_result = scan_result?;
+            let mask = scan_result.full_mask();
             let data = scan_result.raw_data?;
             let record_batch: RecordBatch = data
                 .into_any()
@@ -139,7 +140,7 @@ pub async fn assert_scan_data(engine: Arc<dyn Engine>, test_case: &TestCaseInfo)
             if schema.is_none() {
                 schema = Some(record_batch.schema());
             }
-            if let Some(mask) = scan_result.mask {
+            if let Some(mask) = mask {
                 Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
                 Ok(record_batch)

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

@@ -117,18 +117,14 @@ fn try_main() -> DeltaResult<()> {
         .execute(engine.as_ref())?
         .map(|scan_result| -> DeltaResult<_> {
             let scan_result = scan_result?;
+            let mask = scan_result.full_mask();
             let data = scan_result.raw_data?;
             let record_batch: RecordBatch = data
                 .into_any()
                 .downcast::<ArrowEngineData>()
                 .map_err(|_| delta_kernel::Error::EngineDataType("ArrowEngineData".to_string()))?
                 .into();
-            if let Some(mut mask) = scan_result.mask {
-                let extra_rows = record_batch.num_rows() - mask.len();
-                if extra_rows > 0 {
-                    // we need to extend the mask here in case it's too short
-                    mask.extend(std::iter::repeat(true).take(extra_rows));
-                }
+            if let Some(mask) = mask {
                 Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
                 Ok(record_batch)

kernel/src/scan/mod.rs

@@ -131,7 +131,20 @@ pub struct ScanResult {
     /// plan to call arrow's `filter_record_batch`, you _need_ to extend this vector to the full
     /// length of the batch or arrow will drop the extra rows
     // TODO(nick) this should be allocated by the engine
-    pub mask: Option<Vec<bool>>,
+    raw_mask: Option<Vec<bool>>,
+}
+
+impl ScanResult {
+    /// Returns the raw mask. This is dangerous to use because it may be shorter than expected.
+    pub fn raw_mask(&self) -> Option<&Vec<bool>> {
+        self.raw_mask.as_ref()
+    }
+    /// Extends the underlying mask to match the row count of the accompanying data.
+    pub fn full_mask(&self) -> Option<Vec<bool>> {
+        let mut mask = self.raw_mask.clone()?;
+        mask.resize(self.raw_data.as_ref().ok()?.length(), true);
+        Some(mask)
+    }
 }
 
 /// Scan uses this to set up what kinds of columns it is scanning. For `Selected` we just store the
@@ -314,7 +327,7 @@ impl Scan {
                     let rest = split_vector(sv.as_mut(), len, None);
                     let result = ScanResult {
                         raw_data: logical,
-                        mask: sv,
+                        raw_mask: sv,
                     };
                     selection_vector = rest;
                     Ok(result)

kernel/tests/dv.rs

@@ -16,12 +16,10 @@ fn count_total_scan_rows(
     scan_result_iter
         .map(|scan_result| {
             let scan_result = scan_result?;
-            let data = scan_result.raw_data?;
             // NOTE: The mask only suppresses rows for which it is both present and false.
-            let deleted_rows = scan_result
-                .mask
-                .as_ref()
-                .map_or(0, |mask| mask.iter().filter(|&&m| !m).count());
+            let mask = scan_result.raw_mask();
+            let deleted_rows = mask.into_iter().flatten().filter(|&&m| !m).count();
+            let data = scan_result.raw_data?;
             Ok(data.length() - deleted_rows)
         })
         .fold_ok(0, Add::add)

kernel/tests/golden_tables.rs

@@ -161,14 +161,10 @@ async fn latest_snapshot_test(
     let batches: Vec<RecordBatch> = scan_res
         .map(|scan_result| -> DeltaResult<_> {
             let scan_result = scan_result?;
+            let mask = scan_result.full_mask();
             let data = scan_result.raw_data?;
             let record_batch = to_arrow(data)?;
-            if let Some(mut mask) = scan_result.mask {
-                let extra_rows = record_batch.num_rows() - mask.len();
-                if extra_rows > 0 {
-                    // we need to extend the mask here in case it's too short
-                    mask.extend(std::iter::repeat(true).take(extra_rows));
-                }
+            if let Some(mask) = mask {
                 Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
                 Ok(record_batch)

kernel/tests/read.rs

@@ -401,14 +401,10 @@ fn read_with_execute(
     let batches: Vec<RecordBatch> = scan_results
         .map(|scan_result| -> DeltaResult<_> {
             let scan_result = scan_result?;
+            let mask = scan_result.full_mask();
             let data = scan_result.raw_data?;
             let record_batch = to_arrow(data)?;
-            if let Some(mut mask) = scan_result.mask {
-                let extra_rows = record_batch.num_rows() - mask.len();
-                if extra_rows > 0 {
-                    // we need to extend the mask here in case it's too short
-                    mask.extend(std::iter::repeat(true).take(extra_rows));
-                }
+            if let Some(mask) = mask {
                 Ok(filter_record_batch(&record_batch, &mask.into())?)
             } else {
                 Ok(record_batch)