feat: adds parse options for SQL parser (#3193)

* feat: adds parse options for parser and timezone to scan request

* chore: remove timezone in ScanRequest

* feat: remove timezone in parse options and adds type checking to parititon columns

* fix: comment

* chore: apply suggestions

Co-authored-by: Yingwen <realevenyag@gmail.com>

* fix: format

---------

Co-authored-by: Yingwen <realevenyag@gmail.com>

src/script/src/engine.rs

@@ -20,6 +20,7 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use common_error::ext::ErrorExt;
 use common_query::Output;
+use session::context::{QueryContext, QueryContextRef};
 
 #[async_trait]
 pub trait Script {
@@ -57,8 +58,18 @@ pub trait ScriptEngine {
 }
 
 /// Evaluate script context
-#[derive(Debug, Default)]
-pub struct EvalContext {}
+#[derive(Debug)]
+pub struct EvalContext {
+    pub query_ctx: QueryContextRef,
+}
+
+impl Default for EvalContext {
+    fn default() -> Self {
+        Self {
+            query_ctx: QueryContext::arc(),
+        }
+    }
+}
 
 /// Compile script context
 #[derive(Debug, Default)]

src/script/src/python/engine.rs

@@ -36,7 +36,6 @@ use datatypes::vectors::VectorRef;
 use futures::Stream;
 use query::parser::{QueryLanguageParser, QueryStatement};
 use query::QueryEngineRef;
-use session::context::QueryContextBuilder;
 use snafu::{ensure, ResultExt};
 use sql::statements::statement::Statement;
 
@@ -158,12 +157,16 @@ impl Function for PyUDF {
         let schema = self.fake_schema(columns);
         let columns = columns.to_vec();
         let rb = Some(RecordBatch::new(schema, columns).context(UdfTempRecordBatchSnafu)?);
-        let res = exec_parsed(&self.copr, &rb, &HashMap::new()).map_err(|err| {
-            PyUdfSnafu {
-                msg: format!("{err:#?}"),
-            }
-            .build()
-        })?;
+
+        // FIXME(dennis): Create EvalContext from FunctionContext.
+        let res = exec_parsed(&self.copr, &rb, &HashMap::new(), &EvalContext::default()).map_err(
+            |err| {
+                PyUdfSnafu {
+                    msg: format!("{err:#?}"),
+                }
+                .build()
+            },
+        )?;
         let len = res.columns().len();
         if len == 0 {
             return PyUdfSnafu {
@@ -206,6 +209,7 @@ pub struct CoprStream {
     copr: CoprocessorRef,
     ret_schema: SchemaRef,
     params: HashMap<String, String>,
+    eval_ctx: EvalContext,
 }
 
 impl CoprStream {
@@ -213,6 +217,7 @@ impl CoprStream {
         stream: SendableRecordBatchStream,
         copr: CoprocessorRef,
         params: HashMap<String, String>,
+        eval_ctx: EvalContext,
     ) -> Result<Self> {
         let mut schema = vec![];
         for (ty, name) in copr.return_types.iter().zip(&copr.deco_args.ret_names) {
@@ -238,6 +243,7 @@ impl CoprStream {
             copr,
             ret_schema,
             params,
+            eval_ctx,
         })
     }
 }
@@ -256,9 +262,10 @@ impl Stream for CoprStream {
         match Pin::new(&mut self.stream).poll_next(cx) {
             Poll::Pending => Poll::Pending,
             Poll::Ready(Some(Ok(recordbatch))) => {
-                let batch = exec_parsed(&self.copr, &Some(recordbatch), &self.params)
-                    .map_err(BoxedError::new)
-                    .context(ExternalSnafu)?;
+                let batch =
+                    exec_parsed(&self.copr, &Some(recordbatch), &self.params, &self.eval_ctx)
+                        .map_err(BoxedError::new)
+                        .context(ExternalSnafu)?;
                 Poll::Ready(Some(Ok(batch)))
             }
             Poll::Ready(other) => Poll::Ready(other),
@@ -283,36 +290,35 @@ impl Script for PyScript {
         self
     }
 
-    async fn execute(&self, params: HashMap<String, String>, _ctx: EvalContext) -> Result<Output> {
+    async fn execute(&self, params: HashMap<String, String>, ctx: EvalContext) -> Result<Output> {
         if let Some(sql) = &self.copr.deco_args.sql {
-            let stmt = QueryLanguageParser::parse_sql(sql).unwrap();
+            let stmt = QueryLanguageParser::parse_sql(sql, &ctx.query_ctx).unwrap();
             ensure!(
                 matches!(stmt, QueryStatement::Sql(Statement::Query { .. })),
                 error::UnsupportedSqlSnafu { sql }
             );
-            let ctx = QueryContextBuilder::default().build();
             let plan = self
                 .query_engine
                 .planner()
-                .plan(stmt, ctx.clone())
+                .plan(stmt, ctx.query_ctx.clone())
                 .await
                 .context(DatabaseQuerySnafu)?;
             let res = self
                 .query_engine
-                .execute(plan, ctx)
+                .execute(plan, ctx.query_ctx.clone())
                 .await
                 .context(DatabaseQuerySnafu)?;
             let copr = self.copr.clone();
             match res {
                 Output::Stream(stream) => Ok(Output::Stream(Box::pin(CoprStream::try_new(
-                    stream, copr, params,
+                    stream, copr, params, ctx,
                 )?))),
                 _ => unreachable!(),
             }
         } else {
             let copr = self.copr.clone();
             let params = params.clone();
-            let batch = spawn_blocking_script(move || exec_parsed(&copr, &None, &params))
+            let batch = spawn_blocking_script(move || exec_parsed(&copr, &None, &params, &ctx))
                 .await
                 .context(TokioJoinSnafu)??;
             let batches = RecordBatches::try_new(batch.schema.clone(), vec![batch]).unwrap();

src/script/src/python/ffi_types/copr.rs

@@ -34,12 +34,13 @@ use rustpython_compiler_core::CodeObject;
 use rustpython_vm as vm;
 #[cfg(test)]
 use serde::Deserialize;
-use session::context::QueryContextBuilder;
+use session::context::{QueryContextBuilder, QueryContextRef};
 use snafu::{OptionExt, ResultExt};
 use vm::convert::ToPyObject;
 use vm::{pyclass as rspyclass, PyObjectRef, PyPayload, PyResult, VirtualMachine};
 
 use super::py_recordbatch::PyRecordBatch;
+use crate::engine::EvalContext;
 use crate::python::error::{ensure, ArrowSnafu, OtherSnafu, Result, TypeCastSnafu};
 use crate::python::ffi_types::PyVector;
 #[cfg(feature = "pyo3_backend")]
@@ -325,12 +326,16 @@ pub(crate) fn check_args_anno_real_type(
 /// You can return constant in python code like `return 1, 1.0, True`
 /// which create a constant array(with same value)(currently support int, float and bool) as column on return
 #[cfg(test)]
-pub fn exec_coprocessor(script: &str, rb: &Option<RecordBatch>) -> Result<RecordBatch> {
+pub fn exec_coprocessor(
+    script: &str,
+    rb: &Option<RecordBatch>,
+    eval_ctx: &EvalContext,
+) -> Result<RecordBatch> {
     // 1. parse the script and check if it's only a function with `@coprocessor` decorator, and get `args` and `returns`,
     // 2. also check for exist of `args` in `rb`, if not found, return error
     // cache the result of parse_copr
     let copr = parse::parse_and_compile_copr(script, None)?;
-    exec_parsed(&copr, rb, &HashMap::new())
+    exec_parsed(&copr, rb, &HashMap::new(), eval_ctx)
 }
 
 #[cfg_attr(feature = "pyo3_backend", pyo3class(name = "query_engine"))]
@@ -338,6 +343,7 @@ pub fn exec_coprocessor(script: &str, rb: &Option<RecordBatch>) -> Result<Record
 #[derive(Debug, PyPayload, Clone)]
 pub struct PyQueryEngine {
     inner: QueryEngineWeakRef,
+    query_ctx: QueryContextRef,
 }
 pub(crate) enum Either {
     Rb(RecordBatches),
@@ -365,14 +371,16 @@ impl PyQueryEngine {
 
 #[rspyclass]
 impl PyQueryEngine {
-    pub(crate) fn from_weakref(inner: QueryEngineWeakRef) -> Self {
-        Self { inner }
+    pub(crate) fn from_weakref(inner: QueryEngineWeakRef, query_ctx: QueryContextRef) -> Self {
+        Self { inner, query_ctx }
     }
     pub(crate) fn query_with_new_thread(&self, s: String) -> StdResult<Either, String> {
         let query = self.inner.0.upgrade();
+        let query_ctx = self.query_ctx.clone();
         let thread_handle = std::thread::spawn(move || -> std::result::Result<_, String> {
             if let Some(engine) = query {
-                let stmt = QueryLanguageParser::parse_sql(&s).map_err(|e| e.to_string())?;
+                let stmt =
+                    QueryLanguageParser::parse_sql(&s, &query_ctx).map_err(|e| e.to_string())?;
 
                 // To prevent the error of nested creating Runtime, if is nested, use the parent runtime instead
 
@@ -445,13 +453,14 @@ pub fn exec_parsed(
     copr: &Coprocessor,
     rb: &Option<RecordBatch>,
     params: &HashMap<String, String>,
+    eval_ctx: &EvalContext,
 ) -> Result<RecordBatch> {
     match copr.backend {
-        BackendType::RustPython => rspy_exec_parsed(copr, rb, params),
+        BackendType::RustPython => rspy_exec_parsed(copr, rb, params, eval_ctx),
         BackendType::CPython => {
             #[cfg(feature = "pyo3_backend")]
             {
-                pyo3_exec_parsed(copr, rb, params)
+                pyo3_exec_parsed(copr, rb, params, eval_ctx)
             }
             #[cfg(not(feature = "pyo3_backend"))]
             {
@@ -477,8 +486,9 @@ pub fn exec_copr_print(
     rb: &Option<RecordBatch>,
     ln_offset: usize,
     filename: &str,
+    eval_ctx: &EvalContext,
 ) -> StdResult<RecordBatch, String> {
-    let res = exec_coprocessor(script, rb);
+    let res = exec_coprocessor(script, rb, eval_ctx);
     res.map_err(|e| {
         crate::python::error::pretty_print_error_in_src(script, &e, ln_offset, filename)
     })

src/script/src/python/pyo3/copr_impl.rs

@@ -24,6 +24,7 @@ use pyo3::types::{PyBool, PyDict, PyFloat, PyInt, PyList, PyModule, PyString, Py
 use pyo3::{pymethods, IntoPy, PyAny, PyCell, PyObject, PyResult, Python, ToPyObject};
 use snafu::{ensure, Location, ResultExt};
 
+use crate::engine::EvalContext;
 use crate::python::error::{self, NewRecordBatchSnafu, OtherSnafu, Result};
 use crate::python::ffi_types::copr::PyQueryEngine;
 use crate::python::ffi_types::py_recordbatch::PyRecordBatch;
@@ -65,6 +66,7 @@ pub(crate) fn pyo3_exec_parsed(
     copr: &Coprocessor,
     rb: &Option<RecordBatch>,
     params: &HashMap<String, String>,
+    eval_ctx: &EvalContext,
 ) -> Result<RecordBatch> {
     let _t = metric::METRIC_PYO3_EXEC_TOTAL_ELAPSED.start_timer();
     // i.e params or use `vector(..)` to construct a PyVector
@@ -115,7 +117,7 @@ coprocessor = copr
             let locals = py_main.dict();
 
             if let Some(engine) = &copr.query_engine {
-                let query_engine = PyQueryEngine::from_weakref(engine.clone());
+                let query_engine = PyQueryEngine::from_weakref(engine.clone(), eval_ctx.query_ctx.clone());
                 let query_engine = PyCell::new(py, query_engine)?;
                 globals.set_item("__query__", query_engine)?;
             }
@@ -307,6 +309,7 @@ mod copr_test {
     use datatypes::schema::{ColumnSchema, Schema};
     use datatypes::vectors::{Float32Vector, Float64Vector, VectorRef};
 
+    use super::*;
     use crate::python::ffi_types::copr::{exec_parsed, parse, BackendType};
 
     #[test]
@@ -342,6 +345,7 @@ def a(cpu, mem, **kwargs):
             &copr,
             &Some(rb),
             &HashMap::from([("a".to_string(), "1".to_string())]),
+            &EvalContext::default(),
         );
         dbg!(&ret);
         let _ = ret.unwrap();

src/script/src/python/rspython/copr_impl.rs

@@ -27,6 +27,7 @@ use rustpython_vm::scope::Scope;
 use rustpython_vm::{vm, AsObject, Interpreter, PyObjectRef, PyPayload, VirtualMachine};
 use snafu::{OptionExt, ResultExt};
 
+use crate::engine::EvalContext;
 use crate::python::error::{ensure, ret_other_error_with, NewRecordBatchSnafu, OtherSnafu, Result};
 use crate::python::ffi_types::copr::PyQueryEngine;
 use crate::python::ffi_types::py_recordbatch::PyRecordBatch;
@@ -44,6 +45,7 @@ pub(crate) fn rspy_exec_parsed(
     copr: &Coprocessor,
     rb: &Option<RecordBatch>,
     params: &HashMap<String, String>,
+    eval_ctx: &EvalContext,
 ) -> Result<RecordBatch> {
     let _t = metric::METRIC_RSPY_EXEC_TOTAL_ELAPSED.start_timer();
     // 3. get args from `rb`, and cast them into PyVector
@@ -57,7 +59,7 @@ pub(crate) fn rspy_exec_parsed(
     };
     let interpreter = init_interpreter();
     // 4. then set args in scope and compile then run `CodeObject` which already append a new `Call` node
-    exec_with_cached_vm(copr, rb, args, params, &interpreter)
+    exec_with_cached_vm(copr, rb, args, params, &interpreter, eval_ctx)
 }
 
 /// set arguments with given name and values in python scopes
@@ -100,6 +102,7 @@ pub(crate) fn exec_with_cached_vm(
     args: Vec<PyVector>,
     params: &HashMap<String, String>,
     vm: &Rc<Interpreter>,
+    eval_ctx: &EvalContext,
 ) -> Result<RecordBatch> {
     vm.enter(|vm| -> Result<RecordBatch> {
         let _t = metric::METRIC_RSPY_EXEC_ELAPSED.start_timer();
@@ -116,7 +119,8 @@ pub(crate) fn exec_with_cached_vm(
         }
 
         if let Some(engine) = &copr.query_engine {
-            let query_engine = PyQueryEngine::from_weakref(engine.clone());
+            let query_engine =
+                PyQueryEngine::from_weakref(engine.clone(), eval_ctx.query_ctx.clone());
 
             // put a object named with query of class PyQueryEngine in scope
             set_query_engine_in_scope(&scope, vm, "__query__", query_engine)?;

src/script/src/python/rspython/test.rs

@@ -28,6 +28,7 @@ use ron::from_str as from_ron_string;
 use rustpython_parser::{parse, Mode};
 use serde::{Deserialize, Serialize};
 
+use crate::engine::EvalContext;
 use crate::python::error::{get_error_reason_loc, pretty_print_error_in_src, visualize_loc, Error};
 use crate::python::ffi_types::copr::parse::parse_and_compile_copr;
 use crate::python::ffi_types::copr::{exec_coprocessor, AnnotationInfo};
@@ -124,7 +125,8 @@ fn run_ron_testcases() {
             }
             Predicate::ExecIsOk { fields, columns } => {
                 let rb = create_sample_recordbatch();
-                let res = exec_coprocessor(&testcase.code, &Some(rb)).unwrap();
+                let res =
+                    exec_coprocessor(&testcase.code, &Some(rb), &EvalContext::default()).unwrap();
                 fields
                     .iter()
                     .zip(res.schema.column_schemas())
@@ -150,7 +152,7 @@ fn run_ron_testcases() {
                 reason: part_reason,
             } => {
                 let rb = create_sample_recordbatch();
-                let res = exec_coprocessor(&testcase.code, &Some(rb));
+                let res = exec_coprocessor(&testcase.code, &Some(rb), &EvalContext::default());
                 assert!(res.is_err(), "{res:#?}\nExpect Err(...), actual Ok(...)");
                 if let Err(res) = res {
                     error!(
@@ -252,7 +254,7 @@ def calc_rvs(open_time, close):
         ],
     )
     .unwrap();
-    let ret = exec_coprocessor(python_source, &Some(rb));
+    let ret = exec_coprocessor(python_source, &Some(rb), &EvalContext::default());
     if let Err(Error::PyParse { location: _, error }) = ret {
         let res = visualize_loc(
             python_source,
@@ -298,7 +300,7 @@ def a(cpu, mem):
         ],
     )
     .unwrap();
-    let ret = exec_coprocessor(python_source, &Some(rb));
+    let ret = exec_coprocessor(python_source, &Some(rb), &EvalContext::default());
     if let Err(Error::PyParse { location: _, error }) = ret {
         let res = visualize_loc(
             python_source,