address comments

core/src/execution/datafusion/shuffle_writer.rs

@@ -1282,4 +1282,34 @@ impl RecordBatchStream for EmptyStream {
 }
 
 #[cfg(test)]
-mod shuffle_writer_test;
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_slot_size() {
+        let batch_size = 1usize;
+        // not inclusive of all supported types, but enough to test the function
+        let supported_primitive_types = [
+            DataType::Int32,
+            DataType::Int64,
+            DataType::UInt32,
+            DataType::UInt64,
+            DataType::Float32,
+            DataType::Float64,
+            DataType::Boolean,
+            DataType::Utf8,
+            DataType::LargeUtf8,
+            DataType::Binary,
+            DataType::LargeBinary,
+            DataType::FixedSizeBinary(16),
+        ];
+        let expected_slot_size = [4, 8, 4, 8, 4, 8, 1, 104, 108, 104, 108, 16];
+        supported_primitive_types
+            .iter()
+            .zip(expected_slot_size.iter())
+            .for_each(|(data_type, expected)| {
+                let slot_size = slot_size(batch_size, data_type);
+                assert_eq!(slot_size, *expected);
+            })
+    }
+}