Fix compiler errors in tests

Author: tschwarzinger

datafusion/catalog-listing/src/helpers.rs

@@ -547,6 +547,7 @@ mod tests {
     use datafusion_expr::{
         case, col, lit, AggregateUDF, Expr, LogicalPlan, ScalarUDF, WindowUDF,
     };
+    use datafusion_expr::registry::MemoryExtensionTypeRegistry;
     use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
     use datafusion_physical_plan::ExecutionPlan;
 
@@ -1060,6 +1061,10 @@ mod tests {
             unimplemented!()
         }
 
+        fn extension_types(&self) -> &MemoryExtensionTypeRegistry {
+            unimplemented!()
+        }
+
         fn runtime_env(&self) -> &Arc<RuntimeEnv> {
             unimplemented!()
         }

datafusion/common/src/sort.rs

@@ -1,11 +1,28 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
 use crate::error::{_exec_err, _internal_err};
 use crate::types::SortOrdering;
 use crate::Result;
 use arrow::array::{ArrayRef, DynComparator, UInt32Array};
 use arrow::compute::{partial_sort, SortColumn, SortOptions};
-use std::cmp::Ordering;
 use arrow::datatypes::DataType;
 use arrow::row::{RowConverter, SortField};
+use std::cmp::Ordering;
 
 /// TODO
 #[derive(Clone, Debug, Default, Hash, PartialEq, Eq)]
@@ -175,7 +192,8 @@ pub fn lexsort_to_indices(
         len = limit.min(len);
     }
 
-    let compare_items = columns.iter()
+    let compare_items = columns
+        .iter()
         .map(|c| c.dyn_compartor())
         .collect::<Result<Vec<_>>>()?;
 
@@ -232,7 +250,6 @@ pub(crate) fn lexsort_to_indices_multi_columns(
     Ok(indices)
 }
 
-
 /// we can only do this if the T is primitive
 #[inline]
 fn sort_unstable_by<T, F>(array: &mut [T], limit: usize, cmp: F)

datafusion/core/src/datasource/listing/table.rs

@@ -1199,7 +1199,6 @@ mod tests {
         test::{columns, object_store::register_test_store},
     };
 
-    use arrow::compute::SortOptions;
     use arrow::record_batch::RecordBatch;
     use datafusion_common::stats::Precision;
     use datafusion_common::{assert_contains, ScalarValue};
@@ -1211,6 +1210,8 @@ mod tests {
     use crate::test::object_store::{ensure_head_concurrency, make_test_store_and_state};
     use tempfile::TempDir;
     use url::Url;
+    use datafusion_common::sort::AdvSortOptions;
+    use datafusion_common::types::SortOrdering;
 
     #[tokio::test]
     async fn read_single_file() -> Result<()> {
@@ -1319,7 +1320,8 @@ mod tests {
                 Ok(vec![LexOrdering::new(
                         vec![PhysicalSortExpr {
                             expr: physical_col("string_col", &schema).unwrap(),
-                            options: SortOptions {
+                            options: AdvSortOptions {
+                                ordering: SortOrdering::default(),
                                 descending: false,
                                 nulls_first: false,
                             },

datafusion/core/src/execution/session_state.rs

@@ -1864,7 +1864,10 @@ impl FunctionRegistry for SessionState {
 }
 
 impl ExtensionTypeRegistry for SessionState {
-    fn get_extension_type(&self, name: &str) -> datafusion_common::Result<LogicalTypeRef> {
+    fn get_extension_type(
+        &self,
+        name: &str,
+    ) -> datafusion_common::Result<LogicalTypeRef> {
         self.extension_types.get_extension_type(name)
     }
 

datafusion/core/tests/dataframe/mod.rs

@@ -20,7 +20,13 @@ mod dataframe_functions;
 mod describe;
 mod test_types;
 
-use arrow::array::{as_union_array, record_batch, Array, ArrayRef, AsArray, BooleanArray, DictionaryArray, FixedSizeListArray, FixedSizeListBuilder, Float32Array, Float64Array, Int32Array, Int32Builder, Int8Array, LargeListArray, ListArray, ListBuilder, RecordBatch, StringArray, StringBuilder, StructBuilder, UInt32Array, UInt32Builder, UnionArray, UnionBuilder};
+use arrow::array::{
+    as_union_array, record_batch, Array, ArrayRef, AsArray, BooleanArray,
+    DictionaryArray, FixedSizeListArray, FixedSizeListBuilder, Float32Array,
+    Float64Array, Int32Array, Int32Builder, Int8Array, LargeListArray, ListArray,
+    ListBuilder, RecordBatch, StringArray, StringBuilder, StructBuilder, UInt32Array,
+    UInt32Builder, UnionArray, UnionBuilder,
+};
 use arrow::buffer::ScalarBuffer;
 use arrow::datatypes::{
     DataType, Field, Float32Type, Float64Type, Int32Type, Int64Type, Schema, SchemaRef,
@@ -3111,7 +3117,8 @@ async fn sort_on_union_with_logical_type() -> Result<()> {
         )?),
     )?;
 
-    let record_batch = ctx.table("test_table")
+    let record_batch = ctx
+        .table("test_table")
         .await?
         .sort_by(vec![col("my_union")])?
         .execute_stream()
@@ -3122,8 +3129,14 @@ async fn sort_on_union_with_logical_type() -> Result<()> {
 
     let result = as_union_array(record_batch.column_by_name("my_union").unwrap());
     assert_eq!(result.type_ids(), &[0, 0, 1, 1]);
-    assert_eq!(result.child(0).as_primitive::<Int64Type>().values(), &[-1, 1]);
-    assert_eq!(result.child(1).as_primitive::<Float64Type>().values(), &[3.0, 6.0]);
+    assert_eq!(
+        result.child(0).as_primitive::<Int64Type>().values(),
+        &[-1, 1]
+    );
+    assert_eq!(
+        result.child(1).as_primitive::<Float64Type>().values(),
+        &[3.0, 6.0]
+    );
 
     Ok(())
 }

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -23,7 +23,7 @@ use crate::fuzz_cases::aggregation_fuzzer::{
 };
 
 use arrow::array::{types::Int64Type, Array, ArrayRef, AsArray, Int64Array, RecordBatch};
-use arrow::compute::{concat_batches, SortOptions};
+use arrow::compute::{concat_batches};
 use arrow::datatypes::{
     DataType, IntervalUnit, TimeUnit, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE,
     DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE,
@@ -51,6 +51,7 @@ use test_utils::{add_empty_batches, StringBatchGenerator};
 use rand::rngs::StdRng;
 use rand::{thread_rng, Rng, SeedableRng};
 use tokio::task::JoinSet;
+use datafusion_common::sort::AdvSortOptions;
 
 // ========================================================================
 //  The new aggregation fuzz tests based on [`AggregationFuzzer`]
@@ -315,7 +316,7 @@ async fn run_aggregate_test(input1: Vec<RecordBatch>, group_by_columns: Vec<&str
     for ordering_col in ["a", "b", "c"] {
         sort_keys.push(PhysicalSortExpr {
             expr: col(ordering_col, &schema).unwrap(),
-            options: SortOptions::default(),
+            options: AdvSortOptions::default(),
         })
     }
 

datafusion/core/tests/fuzz_cases/equivalence/ordering.rs

@@ -37,7 +37,8 @@ fn test_ordering_satisfy_with_equivalence_random() -> Result<()> {
     const N_RANDOM_SCHEMA: usize = 5;
     const N_ELEMENTS: usize = 125;
     const N_DISTINCT: usize = 5;
-    const SORT_OPTIONS: SortOptions = SortOptions {
+    const SORT_OPTIONS: AdvSortOptions = AdvSortOptions {
+        ordering: SortOrdering::Default,
         descending: false,
         nulls_first: false,
     };
@@ -63,7 +64,7 @@ fn test_ordering_satisfy_with_equivalence_random() -> Result<()> {
                     .into_iter()
                     .map(|expr| PhysicalSortExpr {
                         expr: Arc::clone(expr),
-                        options: SORT_OPTIONS,
+                        options: SORT_OPTIONS.clone(),
                     })
                     .collect::<LexOrdering>();
                 let expected = is_table_same_after_sort(
@@ -94,7 +95,8 @@ fn test_ordering_satisfy_with_equivalence_complex_random() -> Result<()> {
     const N_RANDOM_SCHEMA: usize = 100;
     const N_ELEMENTS: usize = 125;
     const N_DISTINCT: usize = 5;
-    const SORT_OPTIONS: SortOptions = SortOptions {
+    const SORT_OPTIONS: AdvSortOptions = AdvSortOptions {
+        ordering: SortOrdering::Default,
         descending: false,
         nulls_first: false,
     };
@@ -135,7 +137,7 @@ fn test_ordering_satisfy_with_equivalence_complex_random() -> Result<()> {
                     .into_iter()
                     .map(|expr| PhysicalSortExpr {
                         expr: Arc::clone(expr),
-                        options: SORT_OPTIONS,
+                        options: SORT_OPTIONS.clone(),
                     })
                     .collect::<LexOrdering>();
                 let expected = is_table_same_after_sort(

datafusion/core/tests/fuzz_cases/equivalence/projection.rs

@@ -19,7 +19,6 @@ use crate::fuzz_cases::equivalence::utils::{
     apply_projection, create_random_schema, generate_table_for_eq_properties,
     is_table_same_after_sort, TestScalarUDF,
 };
-use arrow::compute::SortOptions;
 use datafusion_common::Result;
 use datafusion_expr::{Operator, ScalarUDF};
 use datafusion_physical_expr::equivalence::ProjectionMapping;
@@ -29,6 +28,8 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
 use itertools::Itertools;
 use std::sync::Arc;
+use datafusion_common::sort::AdvSortOptions;
+use datafusion_common::types::SortOrdering;
 
 #[test]
 fn project_orderings_random() -> Result<()> {
@@ -108,7 +109,8 @@ fn ordering_satisfy_after_projection_random() -> Result<()> {
     const N_RANDOM_SCHEMA: usize = 20;
     const N_ELEMENTS: usize = 125;
     const N_DISTINCT: usize = 5;
-    const SORT_OPTIONS: SortOptions = SortOptions {
+    const SORT_OPTIONS: AdvSortOptions = AdvSortOptions {
+        ordering: SortOrdering::Default,
         descending: false,
         nulls_first: false,
     };

datafusion/core/tests/fuzz_cases/equivalence/properties.rs

@@ -75,7 +75,7 @@ fn test_find_longest_permutation_random() -> Result<()> {
                     .zip(ordering.iter())
                     .map(|(&idx, sort_expr)| PhysicalSortExpr {
                         expr: Arc::clone(&exprs[idx]),
-                        options: sort_expr.options,
+                        options: sort_expr.options.clone(),
                     })
                     .collect::<LexOrdering>();
                 assert_eq!(

datafusion/core/tests/fuzz_cases/equivalence/utils.rs

@@ -23,7 +23,7 @@ use std::cmp::Ordering;
 use std::sync::Arc;
 
 use arrow::array::{ArrayRef, Float32Array, Float64Array, RecordBatch, UInt32Array};
-use arrow::compute::SortOptions;
+use arrow::compute::{SortColumn, SortOptions};
 use arrow::compute::{lexsort_to_indices, take_record_batch};
 use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
 use datafusion_common::utils::{compare_rows, get_row_at_idx};
@@ -38,6 +38,8 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering;
 
 use itertools::izip;
 use rand::prelude::*;
+use datafusion_common::sort::AdvSortOptions;
+use datafusion_common::types::SortOrdering;
 
 pub fn output_schema(
     mapping: &ProjectionMapping,
@@ -108,7 +110,8 @@ pub fn create_random_schema(seed: u64) -> Result<(SchemaRef, EquivalenceProperti
     let mut rng = StdRng::seed_from_u64(seed);
     let mut remaining_exprs = col_exprs[0..4].to_vec(); // only a, b, c, d are sorted
 
-    let options_asc = SortOptions {
+    let options_asc = AdvSortOptions {
+        ordering: SortOrdering::Default,
         descending: false,
         nulls_first: false,
     };
@@ -121,7 +124,7 @@ pub fn create_random_schema(seed: u64) -> Result<(SchemaRef, EquivalenceProperti
             .drain(0..n_sort_expr)
             .map(|expr| PhysicalSortExpr {
                 expr: Arc::clone(expr),
-                options: options_asc,
+                options: options_asc.clone(),
             })
             .collect();
 
@@ -267,7 +270,7 @@ pub fn is_table_same_after_sort(
             let values = expr_result.into_array(new_batch.num_rows())?;
             Ok(SortColumn {
                 values,
-                options: Some(order_expr.options),
+                options: Some(order_expr.options.to_arrow().unwrap()),
             })
         })
         .collect::<Result<Vec<_>>>()?;
@@ -390,15 +393,15 @@ pub fn generate_table_for_eq_properties(
             .map(
                 |PhysicalSortExpr {
                      expr,
-                     options: options,
+                     options,
                  }| {
                     let col = expr.as_any().downcast_ref::<Column>().unwrap();
                     let (idx, _field) = schema.column_with_name(col.name()).unwrap();
                     let arr = generate_random_array(n_elem, n_distinct);
                     (
                         SortColumn {
                             values: arr,
-                            options: Some(*options),
+                            options: Some(options.to_arrow().unwrap()),
                         },
                         idx,
                     )
@@ -507,7 +510,7 @@ pub fn convert_to_sort_exprs(
         .iter()
         .map(|(expr, options)| PhysicalSortExpr {
             expr: Arc::clone(*expr),
-            options: *options,
+            options: AdvSortOptions::with_default_ordering(*options),
         })
         .collect()
 }
@@ -585,7 +588,7 @@ impl ScalarUDFImpl for TestScalarUDF {
     }
 
     fn output_ordering(&self, input: &[ExprProperties]) -> Result<SortProperties> {
-        Ok(input[0].sort_properties)
+        Ok(input[0].sort_properties.clone())
     }
 
     fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {