Method run

datafusion/datasource/src/file_stream/mod.rs:1441–1551

Runs the test and returns combined stream output and scheduler trace text.

(self)

        /// Runs the test and returns combined stream output and scheduler
        /// trace text.
        async fn run(self) -> Result<String> {
            let observer = self.morselizer.observer().clone();
            observer.clear();

            let metrics_set = ExecutionPlanMetricsSet::new();
            let partition_count = self.num_partitions();

            let mut partitions = (0..partition_count)
                .map(|_| PartitionState::new())
                .collect::<Vec<_>>();

            let mut build_order = Vec::new();
            for partition in self.reads.iter().map(|partition| partition.0) {
                if !build_order.contains(&partition) {
                    build_order.push(partition);
                }
            }
            for partition in 0..partition_count {
                if !build_order.contains(&partition) {
                    build_order.push(partition);
                }
            }

            let config = self.test_config();
            // `DataSourceExec::execute` creates one execution-local shared
            // state object via `create_sibling_state()` and then passes it
            // to `open_with_sibling_state(...)`. These tests build
            // `FileStream`s directly, bypassing `DataSourceExec`, so they must
            // perform the same setup explicitly when exercising sibling-stream
            // work stealing.
            let shared_work_source = config.create_sibling_state().and_then(|state| {
                state.as_ref().downcast_ref::<SharedWorkSource>().cloned()
            });
            if !self.build_streams_on_first_read {
                for partition in build_order {
                    let stream = FileStreamBuilder::new(&config)
                        .with_partition(partition)
                        .with_shared_work_source(shared_work_source.clone())
                        .with_morselizer(Box::new(self.morselizer.clone()))
                        .with_metrics(&metrics_set)
                        .build()?;
                    partitions[partition].set_stream(stream);
                }
            }

            let mut initial_reads: VecDeque<_> = self.reads.into();
            let mut next_round_robin = 0;

            while !initial_reads.is_empty()
                || partitions.iter().any(PartitionState::is_active)
            {
                let partition = if let Some(partition) = initial_reads.pop_front() {
                    partition.0
                } else {
                    let partition = next_round_robin;
                    next_round_robin = (next_round_robin + 1) % partition_count.max(1);
                    partition
                };