fix: correct usage of visitor

Author: feniljain

crates/iceberg/src/arrow/nan_val_cnt_visitor.rs

@@ -1,16 +1,18 @@
 use std::collections::hash_map::Entry;
 use std::collections::HashMap;
+use std::sync::Arc;
 
-use arrow_array::{ArrayRef, Float32Array, Float64Array, RecordBatch};
+use arrow_array::{ArrayRef, Float32Array, Float64Array, RecordBatch, StructArray};
+use arrow_schema::DataType;
 
 use crate::arrow::ArrowArrayAccessor;
 use crate::spec::{
-    visit_schema_with_partner, ListType, MapType, NestedFieldRef, SchemaRef,
-    PrimitiveType, Schema, SchemaWithPartnerVisitor, StructType,
+    visit_struct_with_partner, ListType, MapType, NestedFieldRef, PrimitiveType, Schema, SchemaRef,
+    SchemaWithPartnerVisitor, StructType,
 };
 use crate::Result;
 
-macro_rules! count_float_nans {
+macro_rules! cast_and_update_cnt_map {
     ($t:ty, $col:ident, $self:ident, $field_id:ident) => {
         let nan_val_cnt = $col
             .as_any()
@@ -29,10 +31,24 @@ macro_rules! count_float_nans {
                 v.insert(nan_val_cnt);
             }
         };
+    }
+}
+
+macro_rules! count_float_nans {
+    ($col:ident, $self:ident, $field_id:ident) => {
+        match $col.data_type() {
+            DataType::Float32 => {
+                cast_and_update_cnt_map!(Float32Array, $col, $self, $field_id);
+            }
+            DataType::Float64 => {
+                cast_and_update_cnt_map!(Float64Array, $col, $self, $field_id);
+            }
+            _ => {}
+        }
     };
 }
 
-/// TODO(feniljain)
+/// Visitor which counts and keeps track of NaN value counts in given record batch(s)
 pub struct NanValueCountVisitor {
     /// Stores field ID to NaN value count mapping
     pub nan_value_counts: HashMap<i32, u64>,
@@ -82,21 +98,31 @@ impl SchemaWithPartnerVisitor<ArrayRef> for NanValueCountVisitor {
         Ok(())
     }
 
-    fn primitive(&mut self, p: &PrimitiveType, col: &ArrayRef) -> Result<Self::T> {
-        match p {
-            PrimitiveType::Float => {
-                // let field_id = p.id;
-                // TODO(feniljain): fix this
-                let field_id = 1;
-                count_float_nans!(Float32Array, col, self, field_id);
-            }
-            PrimitiveType::Double => {
-                let field_id = 1;
-                count_float_nans!(Float64Array, col, self, field_id);
-            }
-            _ => {}
-        }
+    fn primitive(&mut self, _p: &PrimitiveType, _col: &ArrayRef) -> Result<Self::T> {
+        Ok(())
+    }
 
+    fn after_struct_field(&mut self, field: &NestedFieldRef, partner: &ArrayRef) -> Result<()> {
+        let field_id = field.id;
+        count_float_nans!(partner, self, field_id);
+        Ok(())
+    }
+
+    fn after_list_element(&mut self, field: &NestedFieldRef, partner: &ArrayRef) -> Result<()> {
+        let field_id = field.id;
+        count_float_nans!(partner, self, field_id);
+        Ok(())
+    }
+
+    fn after_map_key(&mut self, field: &NestedFieldRef, partner: &ArrayRef) -> Result<()> {
+        let field_id = field.id;
+        count_float_nans!(partner, self, field_id);
+        Ok(())
+    }
+
+    fn after_map_value(&mut self, field: &NestedFieldRef, partner: &ArrayRef) -> Result<()> {
+        let field_id = field.id;
+        count_float_nans!(partner, self, field_id);
         Ok(())
     }
 }
@@ -110,14 +136,17 @@ impl NanValueCountVisitor {
     }
 
     /// Compute nan value counts in given schema and record batch
-    pub fn compute(&mut self, schema: SchemaRef, batch: &RecordBatch) -> Result<()> {
-        let arrow_arr_partner_accessor = ArrowArrayAccessor{};
+    pub fn compute(&mut self, schema: SchemaRef, batch: RecordBatch) -> Result<()> {
+        let arrow_arr_partner_accessor = ArrowArrayAccessor {};
 
-        for arr_ref in batch.columns() {
-            visit_schema_with_partner(&schema, arr_ref, self, &arrow_arr_partner_accessor)?;
-        }
+        let struct_arr = Arc::new(StructArray::from(batch)) as ArrayRef;
+        visit_struct_with_partner(
+            &schema.as_struct(),
+            &struct_arr,
+            self,
+            &arrow_arr_partner_accessor,
+        )?;
 
         Ok(())
     }
 }
-

crates/iceberg/src/writer/file_writer/parquet_writer.rs

@@ -525,7 +525,9 @@ impl FileWriter for ParquetWriter {
 
         self.current_row_num += batch.num_rows();
 
-        self.nan_value_count_visitor.compute(self.schema.clone(), batch)?;
+        // TODO(feniljain): Confirm if this `clone` is okay to perform
+        let batch_c = batch.clone();
+        self.nan_value_count_visitor.compute(self.schema.clone(), batch_c)?;
 
         // Lazy initialize the writer
         let writer = if let Some(writer) = &mut self.inner_writer {
@@ -1633,7 +1635,7 @@ mod tests {
             MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
         let file_name_gen =
             DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
-        //
+
         // prepare data
         let arrow_schema = {
             let fields = vec![