Fix DataTypeLayout for LargeList (#3503)

* Fix DataTypeLayout for LargeList * Add datalayout test
apache · Jan 11, 2023 · 55c87c1 · 55c87c1
1 parent ddba53b
commit 55c87c1
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/arrow-array/src/types.rs b/arrow-array/src/types.rs
@@ -767,6 +767,8 @@ pub type LargeBinaryType = GenericBinaryType<i64>;
 #[cfg(test)]
 mod tests {
     use super::*;
+    use arrow_data::{layout, BufferSpec};
+    use std::mem::size_of;
 
     #[test]
     fn month_day_nano_should_roundtrip() {
@@ -803,4 +805,46 @@ mod tests {
         let value = IntervalYearMonthType::make_value(-1, -2);
         assert_eq!(IntervalYearMonthType::to_months(value), -14);
     }
+
+    fn test_layout<T: ArrowPrimitiveType>() {
+        let layout = layout(&T::DATA_TYPE);
+
+        assert_eq!(layout.buffers.len(), 1);
+
+        let spec = &layout.buffers[0];
+        assert_eq!(
+            spec,
+            &BufferSpec::FixedWidth {
+                byte_width: size_of::<T::Native>()
+            }
+        );
+    }
+
+    #[test]
+    fn test_layouts() {
+        test_layout::<Int8Type>();
+        test_layout::<Int16Type>();
+        test_layout::<Int32Type>();
+        test_layout::<Int64Type>();
+        test_layout::<UInt8Type>();
+        test_layout::<UInt16Type>();
+        test_layout::<UInt32Type>();
+        test_layout::<UInt64Type>();
+        test_layout::<Float16Type>();
+        test_layout::<Float32Type>();
+        test_layout::<Float64Type>();
+        test_layout::<TimestampSecondType>();
+        test_layout::<Date32Type>();
+        test_layout::<Date64Type>();
+        test_layout::<Time32SecondType>();
+        test_layout::<Time32MillisecondType>();
+        test_layout::<Time64MicrosecondType>();
+        test_layout::<Time64NanosecondType>();
+        test_layout::<IntervalMonthDayNanoType>();
+        test_layout::<IntervalDayTimeType>();
+        test_layout::<IntervalYearMonthType>();
+        test_layout::<DurationNanosecondType>();
+        test_layout::<DurationMicrosecondType>();
+        test_layout::<DurationMillisecondType>();
+    }
 }
diff --git a/arrow-data/src/data.rs b/arrow-data/src/data.rs
@@ -1470,7 +1470,7 @@ pub fn layout(data_type: &DataType) -> DataTypeLayout {
         DataType::LargeUtf8 => DataTypeLayout::new_binary(size_of::<i64>()),
         DataType::List(_) => DataTypeLayout::new_fixed_width(size_of::<i32>()),
         DataType::FixedSizeList(_, _) => DataTypeLayout::new_empty(), // all in child data
-        DataType::LargeList(_) => DataTypeLayout::new_fixed_width(size_of::<i32>()),
+        DataType::LargeList(_) => DataTypeLayout::new_fixed_width(size_of::<i64>()),
         DataType::Struct(_) => DataTypeLayout::new_empty(), // all in child data,
         DataType::Union(_, _, mode) => {
             let type_ids = BufferSpec::FixedWidth {