local.rs

// 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.

//! [`ExecutionContext`]: DataFusion based execution context for running SQL queries

use std::io::Write;
use std::path::PathBuf;
use std::sync::Arc;

use color_eyre::eyre::eyre;
use datafusion::logical_expr::LogicalPlan;
use futures::TryFutureExt;
use log::{debug, error, info};

use crate::catalog::create_app_catalog;
use crate::config::ExecutionConfig;
use crate::{ExecOptions, ExecResult};
use color_eyre::eyre::{self, Result};
use datafusion::common::Result as DFResult;
use datafusion::execution::{SendableRecordBatchStream, SessionState};
use datafusion::physical_plan::{execute_stream, ExecutionPlan};
use datafusion::prelude::*;
use datafusion::sql::parser::{DFParser, Statement};
use tokio_stream::StreamExt;

use super::executor::dedicated::DedicatedExecutor;
use super::local_benchmarks::{BenchmarkMode, BenchmarkProgressReporter, LocalBenchmarkStats};
use super::stats::{ExecutionDurationStats, ExecutionStats};
#[cfg(feature = "udfs-wasm")]
use super::wasm::create_wasm_udfs;
#[cfg(feature = "observability")]
use {crate::config::ObservabilityConfig, crate::observability::ObservabilityContext};

/// Structure for executing queries
///
/// This context includes both:
///
/// 1. The configuration of a [`SessionContext`] with various extensions enabled
///
/// 2. The code for running SQL queries
///
/// The design goals for this module are to serve as an example of how to integrate
/// DataFusion into an application and to provide a simple interface for running SQL queries
/// with the various extensions enabled.
///
/// Thus it is important (eventually) not depend on the code in the app crate
#[derive(Clone)]
pub struct ExecutionContext {
    config: ExecutionConfig,
    /// Underlying `SessionContext`
    session_ctx: SessionContext,
    /// Path to the configured DDL file
    ddl_path: Option<PathBuf>,
    /// Dedicated executor for running CPU intensive work
    executor: Option<DedicatedExecutor>,
    /// Observability handlers
    #[cfg(feature = "observability")]
    observability: ObservabilityContext,
}

impl std::fmt::Debug for ExecutionContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ExecutionContext").finish()
    }
}

impl ExecutionContext {
    /// Construct a new `ExecutionContext` with the specified configuration
    pub fn try_new(
        config: &ExecutionConfig,
        session_state: SessionState,
        app_name: &str,
        app_version: &str,
    ) -> Result<Self> {
        let mut executor = None;
        if config.dedicated_executor_enabled {
            // Ideally we would only use `enable_time` but we are still doing
            // some network requests as part of planning / execution which require network
            // functionality.
            let runtime_builder = tokio::runtime::Builder::new_multi_thread();
            let dedicated_executor =
                DedicatedExecutor::new("cpu_runtime", config.clone(), runtime_builder);
            executor = Some(dedicated_executor)
        }

        let mut session_ctx = SessionContext::new_with_state(session_state);
        session_ctx = session_ctx.enable_url_table();

        #[cfg(feature = "functions-json")]
        datafusion_functions_json::register_all(&mut session_ctx)?;

        #[cfg(feature = "udfs-wasm")]
        {
            let wasm_udfs = create_wasm_udfs(&config.wasm_udf)?;
            for wasm_udf in wasm_udfs {
                session_ctx.register_udf(wasm_udf);
            }
        }

        session_ctx.register_udtf(
            "parquet_metadata",
            Arc::new(datafusion_functions_parquet::ParquetMetadataFunc {}),
        );
        session_ctx.register_udtf(
            "parquet_page_index",
            Arc::new(datafusion_functions_parquet::ParquetPageIndexFunc {}),
        );

        let catalog = create_app_catalog(config, app_name, app_version)?;
        session_ctx.register_catalog(&config.catalog.name, catalog);

        let ctx = {
            #[cfg(feature = "observability")]
            {
                let observability =
                    ObservabilityContext::try_new(config.observability.clone(), app_name)?;
                if let Some(cat) = session_ctx.catalog(&config.catalog.name) {
                    match cat
                        .register_schema(&config.observability.schema_name, observability.schema())
                    {
                        Ok(_) => {
                            info!("Registered observability schema")
                        }
                        Err(e) => {
                            error!("Error registering observability schema: {}", e)
                        }
                    }
                } else {
                    error!("Missing catalog to register observability schema")
                }
                Self {
                    config: config.clone(),
                    session_ctx,
                    ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
                    executor,
                    observability,
                }
            }
            #[cfg(not(feature = "observability"))]
            {
                Self {
                    config: config.clone(),
                    session_ctx,
                    ddl_path: config.ddl_path.as_ref().map(PathBuf::from),
                    executor,
                }
            }
        };

        Ok(ctx)
    }

    /// Useful for testing execution functionality
    pub fn test() -> Self {
        let cfg = SessionConfig::new().with_information_schema(true);
        let session_ctx = SessionContext::new_with_config(cfg);
        let exec_cfg = ExecutionConfig::default();
        // Okay to `unwrap` in a test
        let app_catalog = create_app_catalog(&exec_cfg, "test", ".0.1.0").unwrap();
        session_ctx.register_catalog("test", app_catalog);
        #[cfg(feature = "observability")]
        let observability =
            ObservabilityContext::try_new(ObservabilityConfig::default(), "test").unwrap();
        Self {
            config: ExecutionConfig::default(),
            session_ctx,
            ddl_path: None,
            executor: None,
            #[cfg(feature = "observability")]
            observability,
        }
    }

    pub fn config(&self) -> &ExecutionConfig {
        &self.config
    }

    pub fn create_tables(&mut self) -> Result<()> {
        Ok(())
    }

    /// Return the inner DataFusion [`SessionContext`]
    pub fn session_ctx(&self) -> &SessionContext {
        &self.session_ctx
    }

    /// Return the inner [`DedicatedExecutor`]
    pub fn executor(&self) -> &Option<DedicatedExecutor> {
        &self.executor
    }

    /// Return the `ObservabilityCtx`
    #[cfg(feature = "observability")]
    pub fn observability(&self) -> &ObservabilityContext {
        &self.observability
    }

    /// Convert the statement to a `LogicalPlan`.  Uses the [`DedicatedExecutor`] if it is available.
    pub async fn statement_to_logical_plan(&self, statement: Statement) -> Result<LogicalPlan> {
        let ctx = self.session_ctx.clone();
        let task = async move { ctx.state().statement_to_plan(statement).await };
        if let Some(executor) = &self.executor {
            let job = executor.spawn(task).map_err(|e| eyre::eyre!(e));
            let job_res = job.await?;
            job_res.map_err(|e| eyre!(e))
        } else {
            task.await.map_err(|e| eyre!(e))
        }
    }

    /// Executes the provided `LogicalPlan` returning a `SendableRecordBatchStream`.  Uses the [`DedicatedExecutor`] if it is available.  Useful on server implementations when planning and execution are done in separate steps and you may be storing the logical plan with something like a request_id.
    pub async fn execute_logical_plan(
        &self,
        logical_plan: LogicalPlan,
    ) -> Result<SendableRecordBatchStream> {
        let ctx = self.session_ctx.clone();
        let task = async move {
            let df = ctx.execute_logical_plan(logical_plan).await?;
            df.execute_stream().await
        };
        if let Some(executor) = &self.executor {
            let job = executor.spawn(task).map_err(|e| eyre!(e));
            let job_res = job.await?;
            job_res.map_err(|e| eyre!(e))
        } else {
            task.await.map_err(|e| eyre!(e))
        }
    }

    /// Executes the specified sql string, driving it to completion but discarding any results
    pub async fn execute_sql_and_discard_results(
        &self,
        sql: &str,
    ) -> datafusion::error::Result<()> {
        let mut stream = self.execute_sql(sql).await?;
        // note we don't call collect() to avoid buffering data
        while let Some(maybe_batch) = stream.next().await {
            maybe_batch?; // check for errors
        }
        Ok(())
    }

    /// Create a physical plan from the specified SQL string.  This is useful if you want to store
    /// the plan and collect metrics from it.
    pub async fn create_physical_plan(
        &self,
        sql: &str,
    ) -> datafusion::error::Result<Arc<dyn ExecutionPlan>> {
        let df = self.session_ctx.sql(sql).await?;
        df.create_physical_plan().await
    }

    /// Executes the specified sql string, returning the resulting
    /// [`SendableRecordBatchStream`] of results
    pub async fn execute_sql(
        &self,
        sql: &str,
    ) -> datafusion::error::Result<SendableRecordBatchStream> {
        self.session_ctx.sql(sql).await?.execute_stream().await
    }

    /// Executes the a pre-parsed DataFusion [`Statement`], returning the
    /// resulting [`SendableRecordBatchStream`] of results
    pub async fn execute_statement(
        &self,
        statement: Statement,
    ) -> datafusion::error::Result<SendableRecordBatchStream> {
        let plan = self
            .session_ctx
            .state()
            .statement_to_plan(statement)
            .await?;
        self.session_ctx
            .execute_logical_plan(plan)
            .await?
            .execute_stream()
            .await
    }

    /// Load DDL from configured DDL path for execution (so strips out comments and empty lines)
    pub fn load_ddl(&self) -> Option<String> {
        info!("Loading DDL from: {:?}", &self.ddl_path);
        if let Some(ddl_path) = &self.ddl_path {
            if ddl_path.exists() {
                let maybe_ddl = std::fs::read_to_string(ddl_path);
                match maybe_ddl {
                    Ok(ddl) => Some(ddl),
                    Err(err) => {
                        error!("Error reading DDL: {:?}", err);
                        None
                    }
                }
            } else {
                info!("DDL path ({:?}) does not exist", ddl_path);
                None
            }
        } else {
            info!("No DDL file configured");
            None
        }
    }

    /// Save DDL to configured DDL path
    pub fn save_ddl(&self, ddl: String) {
        info!("Loading DDL from: {:?}", &self.ddl_path);
        if let Some(ddl_path) = &self.ddl_path {
            match std::fs::File::create(ddl_path) {
                Ok(mut f) => match f.write_all(ddl.as_bytes()) {
                    Ok(_) => {
                        info!("Saved DDL file")
                    }
                    Err(e) => {
                        error!("Error writing DDL file: {e}")
                    }
                },
                Err(e) => {
                    error!("Error creating or opening DDL file: {e}")
                }
            }
        } else {
            info!("No DDL file configured");
        }
    }

    /// Execute DDL statements sequentially
    pub async fn execute_ddl(&self) {
        match self.load_ddl() {
            Some(ddl) => {
                let ddl_statements = ddl.split(';').collect::<Vec<&str>>();
                for statement in ddl_statements {
                    if statement.trim().is_empty() {
                        continue;
                    }
                    if statement.trim().starts_with("--") {
                        continue;
                    }

                    debug!("Executing DDL statement: {:?}", statement);
                    match self.execute_sql_and_discard_results(statement).await {
                        Ok(_) => {
                            info!("DDL statement executed");
                        }
                        Err(e) => {
                            error!("Error executing DDL statement: {e}");
                        }
                    }
                }
            }
            None => {
                info!("No DDL to execute");
            }
        }
    }

    /// Benchmark the provided query.  Currently, only a single statement can be benchmarked
    async fn benchmark_single_iteration(
        &self,
        statement: Statement,
    ) -> Result<(
        usize,
        std::time::Duration,
        std::time::Duration,
        std::time::Duration,
        std::time::Duration,
    )> {
        let start = std::time::Instant::now();
        let logical_plan = self
            .session_ctx()
            .state()
            .statement_to_plan(statement)
            .await?;
        let logical_planning_duration = start.elapsed();
        let physical_plan = self
            .session_ctx()
            .state()
            .create_physical_plan(&logical_plan)
            .await?;
        let physical_planning_duration = start.elapsed();
        let task_ctx = self.session_ctx().task_ctx();
        let mut stream = execute_stream(physical_plan, task_ctx)?;
        let mut rows = 0;
        while let Some(b) = stream.next().await {
            rows += b?.num_rows();
        }
        let execution_duration = start.elapsed();
        let total_duration = start.elapsed();
        Ok((
            rows,
            logical_planning_duration,
            physical_planning_duration - logical_planning_duration,
            execution_duration - physical_planning_duration,
            total_duration,
        ))
    }

    pub async fn benchmark_query(
        &self,
        query: &str,
        cli_iterations: Option<usize>,
        concurrent: bool,
        progress_reporter: Option<Arc<dyn BenchmarkProgressReporter>>,
    ) -> Result<LocalBenchmarkStats> {
        let iterations = cli_iterations.unwrap_or(self.config.benchmark_iterations);
        let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
        let statements = DFParser::parse_sql_with_dialect(query, &dialect)?;

        if statements.len() != 1 {
            return Err(eyre::eyre!("Only a single statement can be benchmarked"));
        }

        let statement = statements[0].clone();
        let concurrency = if concurrent {
            std::cmp::min(iterations, num_cpus::get())
        } else {
            1
        };
        let mode = if concurrent {
            BenchmarkMode::Concurrent(concurrency)
        } else {
            BenchmarkMode::Serial
        };

        info!(
            "Benchmarking query with {} iterations (concurrency: {})",
            iterations, concurrency
        );

        let mut rows_returned = Vec::with_capacity(iterations);
        let mut logical_planning_durations = Vec::with_capacity(iterations);
        let mut physical_planning_durations = Vec::with_capacity(iterations);
        let mut execution_durations = Vec::with_capacity(iterations);
        let mut total_durations = Vec::with_capacity(iterations);

        if !concurrent {
            // Serial execution
            for i in 0..iterations {
                let (rows, lp_dur, pp_dur, exec_dur, total_dur) =
                    self.benchmark_single_iteration(statement.clone()).await?;
                rows_returned.push(rows);
                logical_planning_durations.push(lp_dur);
                physical_planning_durations.push(pp_dur);
                execution_durations.push(exec_dur);
                total_durations.push(total_dur);

                if let Some(ref reporter) = progress_reporter {
                    reporter.on_iteration_complete(i + 1, iterations, total_dur);
                }
            }
        } else {
            // Concurrent execution
            let mut completed = 0;

            while completed < iterations {
                let batch_size = std::cmp::min(concurrency, iterations - completed);
                let mut join_set = tokio::task::JoinSet::new();

                for _ in 0..batch_size {
                    let self_clone = self.clone();
                    let statement_clone = statement.clone();
                    join_set.spawn(async move {
                        self_clone.benchmark_single_iteration(statement_clone).await
                    });
                }

                while let Some(result) = join_set.join_next().await {
                    let (rows, lp_dur, pp_dur, exec_dur, total_dur) = result??;
                    rows_returned.push(rows);
                    logical_planning_durations.push(lp_dur);
                    physical_planning_durations.push(pp_dur);
                    execution_durations.push(exec_dur);
                    total_durations.push(total_dur);

                    completed += 1;
                    if let Some(ref reporter) = progress_reporter {
                        reporter.on_iteration_complete(completed, iterations, total_dur);
                    }
                }
            }
        }

        if let Some(ref reporter) = progress_reporter {
            reporter.finish();
        }

        Ok(LocalBenchmarkStats::new(
            query.to_string(),
            rows_returned,
            mode,
            logical_planning_durations,
            physical_planning_durations,
            execution_durations,
            total_durations,
        ))
    }

    pub async fn analyze_query(&self, query: &str) -> Result<ExecutionStats> {
        let dialect = datafusion::sql::sqlparser::dialect::GenericDialect {};
        let start = std::time::Instant::now();
        let statements = DFParser::parse_sql_with_dialect(query, &dialect)?;
        let parsing_duration = start.elapsed();
        if statements.len() == 1 {
            let statement = statements[0].clone();
            let logical_plan = self
                .session_ctx()
                .state()
                .statement_to_plan(statement.clone())
                .await?;
            let logical_planning_duration = start.elapsed();
            let physical_plan = self
                .session_ctx()
                .state()
                .create_physical_plan(&logical_plan)
                .await?;
            let physical_planning_duration = start.elapsed();
            let task_ctx = self.session_ctx().task_ctx();
            let mut stream = execute_stream(Arc::clone(&physical_plan), task_ctx)?;
            let mut rows = 0;
            let mut batches = 0;
            let mut bytes = 0;
            while let Some(b) = stream.next().await {
                let batch = b?;
                rows += batch.num_rows();
                batches += 1;
                bytes += batch.get_array_memory_size();
            }
            let execution_duration = start.elapsed();
            let durations = ExecutionDurationStats::new(
                parsing_duration,
                logical_planning_duration - parsing_duration,
                physical_planning_duration - logical_planning_duration,
                execution_duration - physical_planning_duration,
                start.elapsed(),
            );
            ExecutionStats::try_new(
                query.to_string(),
                durations,
                rows,
                batches,
                bytes,
                physical_plan,
            )
        } else {
            Err(eyre::eyre!("Only a single statement can be benchmarked"))
        }
    }

    pub async fn execute_sql_with_opts(
        &self,
        sql: &str,
        opts: ExecOptions,
    ) -> DFResult<ExecResult> {
        let df = self.session_ctx.sql(sql).await?;
        let df = if let Some(limit) = opts.limit {
            df.limit(0, Some(limit))?
        } else {
            df
        };
        Ok(ExecResult::RecordBatchStream(df.execute_stream().await?))
    }
}