Method run

datafusion/core/tests/fuzz_cases/sort_fuzz.rs:226–292

Sort the input using SortExec and ensure the results are correct according to `Vec::sort` both with and without spilling

(&self)

    /// Sort the input using SortExec and ensure the results are
    /// correct according to `Vec::sort` both with and without spilling
    async fn run(&self) -> (Vec<Vec<RecordBatch>>, Vec<RecordBatch>) {
        let input = self.input.clone();
        let first_batch = input
            .iter()
            .flat_map(|p| p.iter())
            .next()
            .expect("at least one batch");
        let schema = first_batch.schema();

        let sort_ordering =
            LexOrdering::new(self.sort_columns.iter().map(|c| PhysicalSortExpr {
                expr: col(c, &schema).unwrap(),
                options: SortOptions {
                    descending: false,
                    nulls_first: true,
                },
            }))
            .unwrap();

        let exec = MemorySourceConfig::try_new_exec(&input, schema, None).unwrap();
        let sort = Arc::new(SortExec::new(sort_ordering, exec));

        let session_config = SessionConfig::new().with_repartition_file_scans(false);
        let session_ctx = if let Some(pool_size) = self.pool_size {
            // Make sure there is enough space for the initial spill
            // reservation
            let pool_size = pool_size.saturating_add(
                session_config
                    .options()
                    .execution
                    .sort_spill_reservation_bytes,
            );

            let runtime = RuntimeEnvBuilder::new()
                .with_memory_pool(Arc::new(GreedyMemoryPool::new(pool_size)))
                .build_arc()
                .unwrap();
            SessionContext::new_with_config_rt(session_config, runtime)
        } else {
            SessionContext::new_with_config(session_config)
        };

        let task_ctx = session_ctx.task_ctx();
        let collected = collect(sort.clone(), task_ctx).await.unwrap();

        if self.should_spill {
            assert_ne!(
                sort.metrics().unwrap().spill_count().unwrap(),
                0,
                "Expected spill, but did not: {self:?}"
            );
        } else {
            assert_eq!(
                sort.metrics().unwrap().spill_count().unwrap(),
                0,
                "Expected no spill, but did: {self:?}"
            );
        }