feat: utility functions for creating fsl dtype

apache · Feb 7, 2024 · e38ea82 · e38ea82
1 parent 5572398
commit e38ea82
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 15 deletions.
diff --git a/arrow-array/src/array/list_array.rs b/arrow-array/src/array/list_array.rs
@@ -705,8 +705,7 @@ mod tests {
         let value_offsets = Buffer::from_slice_ref([0i64, 3, 6, 8]);
 
         // Construct a list array from the above two
-        let list_data_type =
-            DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, false)));
+        let list_data_type = DataType::new_large_list(DataType::Int32, false);
         let list_data = ArrayData::builder(list_data_type.clone())
             .len(3)
             .add_buffer(value_offsets.clone())
@@ -863,8 +862,7 @@ mod tests {
         bit_util::set_bit(&mut null_bits, 8);
 
         // Construct a list array from the above two
-        let list_data_type =
-            DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, false)));
+        let list_data_type = DataType::new_large_list(DataType::Int32, false);
         let list_data = ArrayData::builder(list_data_type)
             .len(9)
             .add_buffer(value_offsets)
@@ -929,8 +927,7 @@ mod tests {
         bit_util::set_bit(&mut null_bits, 8);
 
         // Construct a list array from the above two
-        let list_data_type =
-            DataType::LargeList(Arc::new(Field::new("item", DataType::Int32, false)));
+        let list_data_type = DataType::new_large_list(DataType::Int32, false);
         let list_data = ArrayData::builder(list_data_type)
             .len(9)
             .add_buffer(value_offsets)

diff --git a/arrow-ipc/src/convert.rs b/arrow-ipc/src/convert.rs
@@ -877,6 +877,12 @@ mod tests {
                 Field::new("utf8", DataType::Utf8, false),
                 Field::new("binary", DataType::Binary, false),
                 Field::new_list("list[u8]", Field::new("item", DataType::UInt8, false), true),
+                Field::new_fixed_size_list(
+                    "fixed_size_list[u8]",
+                    Field::new("item", DataType::UInt8, false),
+                    2,
+                    true,
+                ),
                 Field::new_list(
                     "list[struct<float32, int32, bool>]",
                     Field::new_struct(

diff --git a/arrow-schema/src/datatype.rs b/arrow-schema/src/datatype.rs
@@ -608,6 +608,24 @@ impl DataType {
     pub fn new_list(data_type: DataType, nullable: bool) -> Self {
         DataType::List(Arc::new(Field::new_list_field(data_type, nullable)))
     }
+
+    /// Create a [`DataType::LargeList`] with elements of the specified type
+    /// and nullability, and conventionally named inner [`Field`] (`"item"`).
+    ///
+    /// To specify field level metadata, construct the inner [`Field`]
+    /// directly via [`Field::new`] or [`Field::new_list_field`].
+    pub fn new_large_list(data_type: DataType, nullable: bool) -> Self {
+        DataType::LargeList(Arc::new(Field::new_list_field(data_type, nullable)))
+    }
+
+    /// Create a [`DataType::FixedSizeList`] with elements of the specified type, size
+    /// and nullability, and conventionally named inner [`Field`] (`"item"`).
+    ///
+    /// To specify field level metadata, construct the inner [`Field`]
+    /// directly via [`Field::new`] or [`Field::new_list_field`].
+    pub fn new_fixed_size_list(data_type: DataType, size: i32, nullable: bool) -> Self {
+        DataType::FixedSizeList(Arc::new(Field::new_list_field(data_type, nullable)), size)
+    }
 }
 
 /// The maximum precision for [DataType::Decimal128] values

diff --git a/arrow-schema/src/field.rs b/arrow-schema/src/field.rs
@@ -217,6 +217,21 @@ impl Field {
         Self::new(name, DataType::LargeList(value.into()), nullable)
     }
 
+    /// Create a new [`Field`] with [`DataType::FixedSizeList`]
+    ///
+    /// - `name`: the name of the [`DataType::FixedSizeList`] field
+    /// - `value`: the description of each list element
+    /// - `size`: the size of the fixed size list
+    /// - `nullable`: if the [`DataType::FixedSizeList`] array is nullable
+    pub fn new_fixed_size_list(
+        name: impl Into<String>,
+        value: impl Into<FieldRef>,
+        size: i32,
+        nullable: bool,
+    ) -> Self {
+        Self::new(name, DataType::FixedSizeList(value.into(), size), nullable)
+    }
+
     /// Create a new [`Field`] with [`DataType::Map`]
     ///
     /// - `name`: the name of the [`DataType::Map`] field

diff --git a/arrow-schema/src/fields.rs b/arrow-schema/src/fields.rs
@@ -444,11 +444,7 @@ mod tests {
                 Field::new("floats", DataType::Struct(floats.clone()), true),
                 true,
             ),
-            Field::new(
-                "f",
-                DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 3),
-                false,
-            ),
+            Field::new_fixed_size_list("f", Field::new("item", DataType::Int32, false), 3, false),
             Field::new_map(
                 "g",
                 "entries",

diff --git a/arrow-select/src/filter.rs b/arrow-select/src/filter.rs
@@ -1261,8 +1261,7 @@ mod tests {
             .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5, 6, 7, 8]))
             .build()
             .unwrap();
-        let list_data_type =
-            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 3);
+        let list_data_type = DataType::new_fixed_size_list(DataType::Int32, 3, false);
         let list_data = ArrayData::builder(list_data_type)
             .len(3)
             .add_child_data(value_data)
@@ -1318,8 +1317,7 @@ mod tests {
         bit_util::set_bit(&mut null_bits, 3);
         bit_util::set_bit(&mut null_bits, 4);
 
-        let list_data_type =
-            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, false)), 2);
+        let list_data_type = DataType::new_fixed_size_list(DataType::Int32, 2, false);
         let list_data = ArrayData::builder(list_data_type)
             .len(5)
             .add_child_data(value_data)