Method run

benchmarks/src/sql_benchmark.rs:204–274

Executes the `run` queries, optionally saving results for later verification. If there are multiple queries only the results for the last query are saved. When `save_results` is `true`, it collects `SELECT`/`WITH` query results and stores them in `last_results`. When `save_results` is `false`, it streams results and counts rows without buffering them. If an 'expect' string is defined this metho

(&mut self, ctx: &SessionContext, save_results: bool)

    /// Returns an error if a `run` query fails or if expected plan strings
    /// are not found.
    pub async fn run(&mut self, ctx: &SessionContext, save_results: bool) -> Result<()> {
        let run_queries = self
            .queries
            .get(&QueryDirective::Run)
            .ok_or_else(|| exec_datafusion_err!("Run query should be loaded by now"))?;

        let mut result_count = 0;

        let result: Vec<RecordBatch> = {
            let mut local_result = vec![];

            for query in run_queries {
                match save_results {
                    true => {
                        debug!(
                            "Running query (saving results) {}-{}: {query}",
                            self.group, self.subgroup
                        );

                        let df = ctx.sql(query).await?;
                        if !self.expect.is_empty() {
                            let physical_plan = df.create_physical_plan().await?;
                            self.validate_expected_plan(&physical_plan)?;
                        }

                        let result_schema = Arc::new(df.schema().as_arrow().clone());
                        let mut batches = df.collect().await?;
                        let trimmed = query.trim_start();

                        // save the output for select/with queries
                        if starts_with_ignore_ascii_case(trimmed, "select")
                            || starts_with_ignore_ascii_case(trimmed, "with")
                        {
                            if batches.is_empty() {
                                batches.push(RecordBatch::new_empty(result_schema));
                            }
                            let row_count_for_query =
                                batches.iter().map(RecordBatch::num_rows).sum::<usize>();
                            debug!(
                                "Persisting {} batches ({} rows)...",
                                batches.len(),
                                row_count_for_query
                            );

                            result_count = row_count_for_query;
                            local_result = batches;
                        }
                    }
                    false => {
                        debug!(
                            "Running query (ignoring results) {}-{}: {query}",
                            self.group, self.subgroup
                        );

                        result_count = self
                            .execute_sql_without_result_buffering(query, ctx)
                            .await?;
                    }