missing testcases

pola-rs · Jan 7, 2025 · df9f62a · df9f62a
1 parent 20622b0
commit df9f62a
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 43 deletions.
diff --git a/crates/polars/tests/it/arrow/array/fixed_size_binary/mod.rs b/crates/polars/tests/it/arrow/array/fixed_size_binary/mod.rs
@@ -1,7 +1,7 @@
 use arrow::array::FixedSizeBinaryArray;
 use arrow::bitmap::Bitmap;
 use arrow::buffer::Buffer;
-use arrow::datatypes::ArrowDataType;
+use arrow::datatypes::{ArrowDataType, ExtensionType};
 
 mod mutable;
 
@@ -94,10 +94,10 @@ fn to() {
     let values = Buffer::from(b"abba".to_vec());
     let a = FixedSizeBinaryArray::new(ArrowDataType::FixedSizeBinary(2), values, None);
 
-    let extension = ArrowDataType::Extension(
-        "a".into(),
-        Box::new(ArrowDataType::FixedSizeBinary(2)),
-        None,
-    );
+    let extension = ArrowDataType::Extension(Box::new(ExtensionType {
+        name: "a".into(),
+        inner: ArrowDataType::FixedSizeBinary(2),
+        metadata: None,
+    }));
     let _ = a.to(extension);
 }
diff --git a/crates/polars/tests/it/arrow/array/growable/list.rs b/crates/polars/tests/it/arrow/array/growable/list.rs
@@ -1,6 +1,6 @@
 use arrow::array::growable::{Growable, GrowableList};
 use arrow::array::{Array, ListArray, MutableListArray, MutablePrimitiveArray, TryExtend};
-use arrow::datatypes::ArrowDataType;
+use arrow::datatypes::{ArrowDataType, ExtensionType};
 
 fn create_list_array(data: Vec<Option<Vec<Option<i32>>>>) -> ListArray<i32> {
     let mut array = MutableListArray::<i32, MutablePrimitiveArray<i32>>::new();
@@ -18,7 +18,11 @@ fn extension() {
 
     let array = create_list_array(data);
 
-    let dtype = ArrowDataType::Extension("ext".into(), Box::new(array.dtype().clone()), None);
+    let dtype = ArrowDataType::Extension(Box::new(ExtensionType {
+        name: "ext".into(),
+        inner: array.dtype().clone(),
+        metadata: None,
+    }));
     let array_ext = ListArray::new(
         dtype,
         array.offsets().clone(),

diff --git a/crates/polars/tests/it/arrow/array/growable/mod.rs b/crates/polars/tests/it/arrow/array/growable/mod.rs
@@ -11,7 +11,7 @@ mod utf8;
 
 use arrow::array::growable::make_growable;
 use arrow::array::*;
-use arrow::datatypes::{ArrowDataType, Field};
+use arrow::datatypes::{ArrowDataType, ExtensionType, Field};
 
 #[test]
 fn test_make_growable() {
@@ -44,20 +44,20 @@ fn test_make_growable_extension() {
     .unwrap();
     make_growable(&[&array], false, 2);
 
-    let dtype = ArrowDataType::Extension("ext".into(), Box::new(ArrowDataType::Int32), None);
+    let dtype = ArrowDataType::Extension(Box::new(ExtensionType {
+        name: "ext".into(),
+        inner: ArrowDataType::Int32,
+        metadata: None,
+    }));
     let array = Int32Array::from_slice([1, 2]).to(dtype.clone());
     let array_grown = make_growable(&[&array], false, 2).as_box();
     assert_eq!(array_grown.dtype(), &dtype);
 
-    let dtype = ArrowDataType::Extension(
-        "ext".into(),
-        Box::new(ArrowDataType::Struct(vec![Field::new(
-            "a".into(),
-            ArrowDataType::Int32,
-            false,
-        )])),
-        None,
-    );
+    let dtype = ArrowDataType::Extension(Box::new(ExtensionType {
+        name: "ext".into(),
+        inner: ArrowDataType::Struct(vec![Field::new("a".into(), ArrowDataType::Int32, false)]),
+        metadata: None,
+    }));
     let array = StructArray::new(
         dtype.clone(),
         2,

diff --git a/crates/polars/tests/it/arrow/array/mod.rs b/crates/polars/tests/it/arrow/array/mod.rs
@@ -15,7 +15,8 @@ mod utf8;
 
 use arrow::array::{clone, new_empty_array, new_null_array, Array, PrimitiveArray};
 use arrow::bitmap::Bitmap;
-use arrow::datatypes::{ArrowDataType, Field, UnionMode};
+use arrow::datatypes::{ArrowDataType, ExtensionType, Field, UnionMode};
+use union::union_type;
 
 #[test]
 fn nulls() {
@@ -37,12 +38,12 @@ fn nulls() {
 
     // unions' null count is always 0
     let datatypes = vec![
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Dense,
         ),
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Sparse,
@@ -68,15 +69,19 @@ fn empty() {
         ))),
         ArrowDataType::List(Box::new(Field::new(
             "a".into(),
-            ArrowDataType::Extension("ext".into(), Box::new(ArrowDataType::Int32), None),
+            ArrowDataType::Extension(Box::new(ExtensionType {
+                name: "ext".into(),
+                inner: ArrowDataType::Int32,
+                metadata: None,
+            })),
             true,
         ))),
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Sparse,
         ),
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Dense,
@@ -99,12 +104,12 @@ fn empty_extension() {
             ArrowDataType::Binary,
             true,
         ))),
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Sparse,
         ),
-        ArrowDataType::Union(
+        union_type(
             vec![Field::new("a".into(), ArrowDataType::Binary, true)],
             None,
             UnionMode::Dense,
@@ -113,10 +118,16 @@ fn empty_extension() {
     ];
     let a = datatypes
         .into_iter()
-        .map(|dt| ArrowDataType::Extension("ext".into(), Box::new(dt), None))
+        .map(|dt| {
+            ArrowDataType::Extension(Box::new(ExtensionType {
+                name: "ext".into(),
+                inner: dt,
+                metadata: None,
+            }))
+        })
         .all(|x| {
             let a = new_empty_array(x);
-            a.len() == 0 && matches!(a.dtype(), ArrowDataType::Extension(_, _, _))
+            a.len() == 0 && matches!(a.dtype(), ArrowDataType::Extension(_))
         });
     assert!(a);
 }

diff --git a/crates/polars/tests/it/arrow/array/union.rs b/crates/polars/tests/it/arrow/array/union.rs
@@ -4,6 +4,10 @@ use arrow::datatypes::*;
 use arrow::scalar::{new_scalar, PrimitiveScalar, Scalar, UnionScalar, Utf8Scalar};
 use polars_error::PolarsResult;
 
+pub fn union_type(fields: Vec<Field>, ids: Option<Vec<i32>>, mode: UnionMode) -> ArrowDataType {
+    ArrowDataType::Union(Box::new(UnionType { fields, ids, mode }))
+}
+
 fn next_unwrap<T, I>(iter: &mut I) -> T
 where
     I: Iterator<Item = Box<dyn Scalar>>,
@@ -23,7 +27,7 @@ fn sparse_debug() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let types = vec![0, 0, 1].into();
     let fields = vec![
         Int32Array::from(&[Some(1), None, Some(2)]).boxed(),
@@ -43,7 +47,7 @@ fn dense_debug() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Dense);
+    let dtype = union_type(fields, None, UnionMode::Dense);
     let types = vec![0, 0, 1].into();
     let fields = vec![
         Int32Array::from(&[Some(1), None, Some(2)]).boxed(),
@@ -64,7 +68,7 @@ fn slice() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::LargeUtf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let types = Buffer::from(vec![0, 0, 1]);
     let fields = vec![
         Int32Array::from(&[Some(1), None, Some(2)]).boxed(),
@@ -92,7 +96,7 @@ fn iter_sparse() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let types = Buffer::from(vec![0, 0, 1]);
     let fields = vec![
         Int32Array::from(&[Some(1), None, Some(2)]).boxed(),
@@ -125,7 +129,7 @@ fn iter_dense() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Dense);
+    let dtype = union_type(fields, None, UnionMode::Dense);
     let types = Buffer::from(vec![0, 0, 1]);
     let offsets = Buffer::<i32>::from(vec![0, 1, 0]);
     let fields = vec![
@@ -159,7 +163,7 @@ fn iter_sparse_slice() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let types = Buffer::from(vec![0, 0, 1]);
     let fields = vec![
         Int32Array::from(&[Some(1), Some(3), Some(2)]).boxed(),
@@ -185,7 +189,7 @@ fn iter_dense_slice() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Dense);
+    let dtype = union_type(fields, None, UnionMode::Dense);
     let types = Buffer::from(vec![0, 0, 1]);
     let offsets = Buffer::<i32>::from(vec![0, 1, 0]);
     let fields = vec![
@@ -212,7 +216,7 @@ fn scalar() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Dense);
+    let dtype = union_type(fields, None, UnionMode::Dense);
     let types = Buffer::from(vec![0, 0, 1]);
     let offsets = Buffer::<i32>::from(vec![0, 1, 0]);
     let fields = vec![
@@ -269,7 +273,7 @@ fn dense_without_offsets_is_error() {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Dense);
+    let dtype = union_type(fields, None, UnionMode::Dense);
     let types = vec![0, 0, 1].into();
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
@@ -285,7 +289,7 @@ fn fields_must_match() {
         Field::new("a".into(), ArrowDataType::Int64, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let types = vec![0, 0, 1].into();
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
@@ -301,7 +305,7 @@ fn sparse_with_offsets_is_error() {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
         Utf8Array::<i32>::from([Some("a"), Some("b"), Some("c")]).boxed(),
@@ -319,7 +323,7 @@ fn offsets_must_be_in_bounds() {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
         Utf8Array::<i32>::from([Some("a"), Some("b"), Some("c")]).boxed(),
@@ -338,7 +342,7 @@ fn sparse_with_wrong_offsets1_is_error() {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
         Utf8Array::<i32>::from([Some("a"), Some("b"), Some("c")]).boxed(),
@@ -357,7 +361,7 @@ fn types_must_be_in_bounds() -> PolarsResult<()> {
         Field::new("a".into(), ArrowDataType::Int32, true),
         Field::new("b".into(), ArrowDataType::Utf8, true),
     ];
-    let dtype = ArrowDataType::Union(fields, None, UnionMode::Sparse);
+    let dtype = union_type(fields, None, UnionMode::Sparse);
     let fields = vec![
         Int32Array::from([Some(1), Some(3), Some(2)]).boxed(),
         Utf8Array::<i32>::from([Some("a"), Some("b"), Some("c")]).boxed(),