Cache logical nulls of the array to avoid clones or repeated recomputation

Author: jhorstmann

parquet/src/arrow/arrow_writer/levels.rs

@@ -498,10 +498,10 @@ impl LevelInfoBuilder {
                 def_levels.reserve(len);
                 info.non_null_indices.reserve(len);
 
-                match info.array.logical_nulls() {
+                match &info.logical_nulls {
                     Some(nulls) => {
                         assert!(range.end <= nulls.len());
-                        let nulls = nulls.into_inner();
+                        let nulls = nulls.inner();
                         def_levels.extend(range.clone().map(|i| {
                             // Safety: range.end was asserted to be in bounds earlier
                             let valid = unsafe { nulls.value_unchecked(i) };
@@ -568,6 +568,9 @@ pub(crate) struct ArrayLevels {
 
     /// The arrow array
     array: ArrayRef,
+
+    /// cached logical nulls of the array.
+    logical_nulls: Option<NullBuffer>,
 }
 
 impl PartialEq for ArrayLevels {
@@ -578,6 +581,7 @@ impl PartialEq for ArrayLevels {
             && self.max_def_level == other.max_def_level
             && self.max_rep_level == other.max_rep_level
             && self.array.as_ref() == other.array.as_ref()
+            && self.logical_nulls.as_ref() == other.logical_nulls.as_ref()
     }
 }
 impl Eq for ArrayLevels {}
@@ -590,13 +594,16 @@ impl ArrayLevels {
             false => ctx.def_level,
         };
 
+        let logical_nulls = array.logical_nulls();
+
         Self {
             def_levels: (max_def_level != 0).then(Vec::new),
             rep_levels: (max_rep_level != 0).then(Vec::new),
             non_null_indices: vec![],
             max_def_level,
             max_rep_level,
             array,
+            logical_nulls,
         }
     }
 
@@ -670,6 +677,7 @@ mod tests {
             max_def_level: 2,
             max_rep_level: 2,
             array: Arc::new(primitives),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected);
     }
@@ -690,6 +698,7 @@ mod tests {
             max_def_level: 0,
             max_rep_level: 0,
             array,
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
     }
@@ -709,13 +718,15 @@ mod tests {
         let levels = calculate_array_levels(&array, &field).unwrap();
         assert_eq!(levels.len(), 1);
 
+        let logical_nulls = array.logical_nulls();
         let expected_levels = ArrayLevels {
             def_levels: Some(vec![1, 0, 1, 1, 0]),
             rep_levels: None,
             non_null_indices: vec![0, 2, 3],
             max_def_level: 1,
             max_rep_level: 0,
             array,
+            logical_nulls,
         };
         assert_eq!(&levels[0], &expected_levels);
     }
@@ -750,6 +761,7 @@ mod tests {
             max_def_level: 1,
             max_rep_level: 1,
             array: Arc::new(leaf_array),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
 
@@ -783,6 +795,7 @@ mod tests {
             max_def_level: 2,
             max_rep_level: 1,
             array: Arc::new(leaf_array),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
     }
@@ -832,6 +845,7 @@ mod tests {
             max_def_level: 3,
             max_rep_level: 1,
             array: Arc::new(leaf),
+            logical_nulls: None,
         };
 
         assert_eq!(&levels[0], &expected_levels);
@@ -882,6 +896,7 @@ mod tests {
             max_def_level: 5,
             max_rep_level: 2,
             array: Arc::new(leaf),
+            logical_nulls: None,
         };
 
         assert_eq!(&levels[0], &expected_levels);
@@ -919,6 +934,7 @@ mod tests {
             max_def_level: 1,
             max_rep_level: 1,
             array: Arc::new(leaf),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
 
@@ -951,6 +967,7 @@ mod tests {
             max_def_level: 3,
             max_rep_level: 1,
             array: Arc::new(leaf),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
 
@@ -999,6 +1016,7 @@ mod tests {
             max_def_level: 5,
             max_rep_level: 2,
             array: Arc::new(leaf),
+            logical_nulls: None,
         };
         assert_eq!(&levels[0], &expected_levels);
     }
@@ -1031,13 +1049,15 @@ mod tests {
         let levels = calculate_array_levels(&a_array, &a_field).unwrap();
         assert_eq!(levels.len(), 1);
 
+        let logical_nulls = leaf.logical_nulls();
         let expected_levels = ArrayLevels {
             def_levels: Some(vec![3, 2, 3, 1, 0, 3]),
             rep_levels: None,
             non_null_indices: vec![0, 2, 5],
             max_def_level: 3,
             max_rep_level: 0,
             array: leaf,
+            logical_nulls,
         };
         assert_eq!(&levels[0], &expected_levels);
     }
@@ -1077,6 +1097,7 @@ mod tests {
             max_def_level: 3,
             max_rep_level: 1,
             array: Arc::new(a_values),
+            logical_nulls: None,
         };
         assert_eq!(list_level, &expected_level);
     }
@@ -1169,45 +1190,52 @@ mod tests {
             max_def_level: 0,
             max_rep_level: 0,
             array: Arc::new(a),
+            logical_nulls: None,
         };
         assert_eq!(list_level, &expected_level);
 
         // test "b" levels
         let list_level = levels.get(1).unwrap();
 
+        let b_logical_nulls = b.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![1, 0, 0, 1, 1]),
             rep_levels: None,
             non_null_indices: vec![0, 3, 4],
             max_def_level: 1,
             max_rep_level: 0,
             array: Arc::new(b),
+            logical_nulls: b_logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
 
         // test "d" levels
         let list_level = levels.get(2).unwrap();
 
+        let d_logical_nulls = d.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![1, 1, 1, 2, 1]),
             rep_levels: None,
             non_null_indices: vec![3],
             max_def_level: 2,
             max_rep_level: 0,
             array: Arc::new(d),
+            logical_nulls: d_logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
 
         // test "f" levels
         let list_level = levels.get(3).unwrap();
 
+        let f_logical_nulls = f.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![3, 2, 3, 2, 3]),
             rep_levels: None,
             non_null_indices: vec![0, 2, 4],
             max_def_level: 3,
             max_rep_level: 0,
             array: Arc::new(f),
+            logical_nulls: f_logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
     }
@@ -1303,6 +1331,7 @@ mod tests {
         assert_eq!(levels.len(), 2);
 
         let map = batch.column(0).as_map();
+        let map_keys_logical_nulls = map.keys().logical_nulls();
 
         // test key levels
         let list_level = &levels[0];
@@ -1314,11 +1343,13 @@ mod tests {
             max_def_level: 1,
             max_rep_level: 1,
             array: map.keys().clone(),
+            logical_nulls: map_keys_logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
 
         // test values levels
         let list_level = levels.get(1).unwrap();
+        let map_values_logical_nulls = map.values().logical_nulls();
 
         let expected_level = ArrayLevels {
             def_levels: Some(vec![2, 2, 2, 1, 2, 1, 2]),
@@ -1327,6 +1358,7 @@ mod tests {
             max_def_level: 2,
             max_rep_level: 1,
             array: map.values().clone(),
+            logical_nulls: map_values_logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
     }
@@ -1405,13 +1437,15 @@ mod tests {
         let levels = calculate_array_levels(rb.column(0), rb.schema().field(0)).unwrap();
         let list_level = &levels[0];
 
+        let logical_nulls = values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![4, 1, 0, 2, 2, 3, 4]),
             rep_levels: Some(vec![0, 0, 0, 0, 1, 0, 0]),
             non_null_indices: vec![0, 4],
             max_def_level: 4,
             max_rep_level: 1,
             array: values,
+            logical_nulls,
         };
 
         assert_eq!(list_level, &expected_level);
@@ -1445,13 +1479,15 @@ mod tests {
 
         assert_eq!(levels.len(), 1);
 
+        let logical_nulls = values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![4, 4, 3, 2, 0, 4, 4, 0, 1]),
             rep_levels: Some(vec![0, 1, 0, 0, 0, 0, 1, 0, 0]),
             non_null_indices: vec![0, 1, 5, 6],
             max_def_level: 4,
             max_rep_level: 1,
             array: values,
+            logical_nulls,
         };
 
         assert_eq!(&levels[0], &expected_level);
@@ -1530,24 +1566,28 @@ mod tests {
 
         assert_eq!(levels.len(), 2);
 
+        let a1_logical_nulls = a1_values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![0, 0, 1, 6, 5, 2, 3, 1]),
             rep_levels: Some(vec![0, 0, 0, 0, 2, 0, 1, 0]),
             non_null_indices: vec![1],
             max_def_level: 6,
             max_rep_level: 2,
             array: a1_values,
+            logical_nulls: a1_logical_nulls,
         };
 
         assert_eq!(&levels[0], &expected_level);
 
+        let a2_logical_nulls = a2_values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![0, 0, 1, 3, 2, 4, 1]),
             rep_levels: Some(vec![0, 0, 0, 0, 0, 1, 0]),
             non_null_indices: vec![4],
             max_def_level: 4,
             max_rep_level: 1,
             array: a2_values,
+            logical_nulls: a2_logical_nulls,
         };
 
         assert_eq!(&levels[1], &expected_level);
@@ -1579,13 +1619,15 @@ mod tests {
 
         let list_level = &levels[0];
 
+        let logical_nulls = values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![0, 0, 3, 3]),
             rep_levels: Some(vec![0, 0, 0, 1]),
             non_null_indices: vec![6, 7],
             max_def_level: 3,
             max_rep_level: 1,
             array: values,
+            logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
     }
@@ -1729,22 +1771,26 @@ mod tests {
         let b_levels = &levels[1];
 
         // [[{a: 1}, null], null, [null, null], [{a: null}, {a: 2}]]
+        let values_a_logical_nulls = values_a.logical_nulls();
         let expected_a = ArrayLevels {
             def_levels: Some(vec![4, 2, 0, 2, 2, 3, 4]),
             rep_levels: Some(vec![0, 1, 0, 0, 1, 0, 1]),
             non_null_indices: vec![0, 7],
             max_def_level: 4,
             max_rep_level: 1,
             array: values_a,
+            logical_nulls: values_a_logical_nulls,
         };
         // [[{b: 2}, null], null, [null, null], [{b: 3}, {b: 4}]]
+        let values_b_logical_nulls = values_b.logical_nulls();
         let expected_b = ArrayLevels {
             def_levels: Some(vec![3, 2, 0, 2, 2, 3, 3]),
             rep_levels: Some(vec![0, 1, 0, 0, 1, 0, 1]),
             non_null_indices: vec![0, 6, 7],
             max_def_level: 3,
             max_rep_level: 1,
             array: values_b,
+            logical_nulls: values_b_logical_nulls,
         };
 
         assert_eq!(a_levels, &expected_a);
@@ -1769,13 +1815,15 @@ mod tests {
 
         let list_level = &levels[0];
 
+        let logical_nulls = values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![1, 0, 1]),
             rep_levels: Some(vec![0, 0, 0]),
             non_null_indices: vec![],
             max_def_level: 3,
             max_rep_level: 1,
             array: values,
+            logical_nulls,
         };
         assert_eq!(list_level, &expected_level);
     }
@@ -1804,13 +1852,15 @@ mod tests {
         builder.write(0..4);
         let levels = builder.finish();
 
+        let logical_nulls = values.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![5, 4, 5, 2, 5, 3, 5, 5, 4, 4, 0]),
             rep_levels: Some(vec![0, 2, 2, 1, 0, 1, 0, 2, 1, 2, 0]),
             non_null_indices: vec![0, 2, 3, 4, 5],
             max_def_level: 5,
             max_rep_level: 2,
             array: values,
+            logical_nulls,
         };
 
         assert_eq!(levels[0], expected_level);
@@ -1834,13 +1884,16 @@ mod tests {
         let mut builder = levels(&item_field, dict.clone());
         builder.write(0..4);
         let levels = builder.finish();
+
+        let logical_nulls = dict.logical_nulls();
         let expected_level = ArrayLevels {
             def_levels: Some(vec![0, 0, 1, 1]),
             rep_levels: None,
             non_null_indices: vec![2, 3],
             max_def_level: 1,
             max_rep_level: 0,
             array: Arc::new(dict),
+            logical_nulls,
         };
         assert_eq!(levels[0], expected_level);
     }