add two more unit test cases

apache · Jun 8, 2021 · cfa25b3 · cfa25b3
1 parent 1e4c160
commit cfa25b3
Showing 1 changed file with 50 additions and 0 deletions.
diff --git a/arrow/src/compute/kernels/partition.rs b/arrow/src/compute/kernels/partition.rs
@@ -100,6 +100,7 @@ mod tests {
     use super::*;
     use crate::array::*;
     use crate::compute::SortOptions;
+    use crate::datatypes::DataType;
     use std::sync::Arc;
 
     #[test]
@@ -153,6 +154,55 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_lexicographical_partition_all_equal_values() -> Result<()> {
+        let input = vec![SortColumn {
+            values: Arc::new(Int64Array::from_value(1, 1000)) as ArrayRef,
+            options: Some(SortOptions {
+                descending: false,
+                nulls_first: true,
+            }),
+        }];
+        {
+            let results = lexicographical_partition_points(&input)?;
+            assert_eq!(vec![0, 1000], results);
+        }
+        {
+            let results = lexicographical_partition_ranges(&input)?;
+            assert_eq!(vec![(0_usize..1000_usize)], results);
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_lexicographical_partition_all_null_values() -> Result<()> {
+        let input = vec![
+            SortColumn {
+                values: new_null_array(&DataType::Int8, 1000),
+                options: Some(SortOptions {
+                    descending: false,
+                    nulls_first: true,
+                }),
+            },
+            SortColumn {
+                values: new_null_array(&DataType::UInt16, 1000),
+                options: Some(SortOptions {
+                    descending: false,
+                    nulls_first: false,
+                }),
+            },
+        ];
+        {
+            let results = lexicographical_partition_points(&input)?;
+            assert_eq!(vec![0, 1000], results);
+        }
+        {
+            let results = lexicographical_partition_ranges(&input)?;
+            assert_eq!(vec![(0_usize..1000_usize)], results);
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_lexicographical_partition_unique_column_1() -> Result<()> {
         let input = vec![