feat(expr): distributed make timestamptz

e2e_test/batch/basic/local/make_timestamptz.slt.part → e2e_test/batch/basic/make_timestamptz.slt.part

@@ -31,6 +31,20 @@ SELECT make_timestamptz(1973, 07, 15, 08, 15, 55.33);
 ----
 1973-07-15 08:15:55.330-04:00
 
+statement ok
+create table ttz(tstz timestamptz);
+
+statement ok
+insert into ttz values(make_timestamptz(1973, 06, 25, 08, 15, 55.33));
+
+query TT
+select * from ttz;
+----
+1973-06-25 08:15:55.330-04:00
+
+statement ok
+drop table ttz;
+
 query error Invalid parameter time_zone: 'Nehwon/Lankhmar' is not a valid timezone
 SELECT make_timestamptz(1910, 12, 24, 0, 0, 0, 'Nehwon/Lankhmar');
 

proto/stream_plan.proto

@@ -878,3 +878,8 @@ message StreamFragmentGraph {
   // If none, default parallelism will be applied.
   Parallelism parallelism = 6;
 }
+
+// Provide statement-local context, e.g. session info like time zone, for runtime execution.
+message CapturedContext {
+  string time_zone = 1;
+}

proto/task_service.proto

@@ -40,6 +40,7 @@ message CreateTaskRequest {
   batch_plan.PlanFragment plan = 2;
   common.BatchQueryEpoch epoch = 3;
   map<string, string> tracing_context = 4;
+  stream_plan.CapturedContext captured_context = 5;
 }
 
 message CancelTaskRequest {
@@ -63,6 +64,7 @@ message ExecuteRequest {
   batch_plan.PlanFragment plan = 2;
   common.BatchQueryEpoch epoch = 3;
   map<string, string> tracing_context = 4;
+  stream_plan.CapturedContext captured_context = 5;
 }
 
 service TaskService {

src/batch/src/execution/grpc_exchange.rs

@@ -16,6 +16,7 @@ use std::fmt::{Debug, Formatter};
 
 use futures::StreamExt;
 use risingwave_common::array::DataChunk;
+use risingwave_expr::captured_context::capture_context;
 use risingwave_pb::batch_plan::exchange_source::LocalExecutePlan::{self, Plan};
 use risingwave_pb::batch_plan::TaskOutputId;
 use risingwave_pb::task_service::{ExecuteRequest, GetDataResponse};
@@ -50,6 +51,7 @@ impl GrpcExchangeSource {
                     plan: plan.plan,
                     epoch: plan.epoch,
                     tracing_context: plan.tracing_context,
+                    captured_context: Some(capture_context()?),
                 };
                 client.execute(execute_request).await?
             }

src/batch/src/rpc/service/task_service.rs

@@ -63,6 +63,7 @@ impl TaskService for BatchServiceImpl {
             plan,
             epoch,
             tracing_context,
+            captured_context,
         } = request.into_inner();
 
         let (state_tx, state_rx) = tokio::sync::mpsc::channel(TASK_STATUS_BUFFER_SIZE);
@@ -79,6 +80,7 @@ impl TaskService for BatchServiceImpl {
                 ),
                 state_reporter,
                 TracingContext::from_protobuf(&tracing_context),
+                captured_context.expect("no captured context found"),
             )
             .await;
         match res {
@@ -119,12 +121,14 @@ impl TaskService for BatchServiceImpl {
             plan,
             epoch,
             tracing_context,
+            captured_context,
         } = req.into_inner();
 
         let task_id = task_id.expect("no task id found");
         let plan = plan.expect("no plan found").clone();
         let epoch = epoch.expect("no epoch found");
         let tracing_context = TracingContext::from_protobuf(&tracing_context);
+        let captured_context = captured_context.expect("no captured context found");
 
         let context = ComputeNodeContext::new_for_local(self.env.clone());
         trace!(
@@ -135,7 +139,11 @@ impl TaskService for BatchServiceImpl {
         let task = BatchTaskExecution::new(&task_id, plan, context, epoch, self.mgr.runtime())?;
         let task = Arc::new(task);
         let (tx, rx) = tokio::sync::mpsc::channel(LOCAL_EXECUTE_BUFFER_SIZE);
-        if let Err(e) = task.clone().async_execute(None, tracing_context).await {
+        if let Err(e) = task
+            .clone()
+            .async_execute(None, tracing_context, captured_context)
+            .await
+        {
             error!(
                 "failed to build executors and trigger execution of Task {:?}: {}",
                 task_id, e

src/batch/src/task/task_execution.rs

@@ -27,8 +27,10 @@ use risingwave_common::array::DataChunk;
 use risingwave_common::util::panic::FutureCatchUnwindExt;
 use risingwave_common::util::runtime::BackgroundShutdownRuntime;
 use risingwave_common::util::tracing::TracingContext;
+use risingwave_expr::captured_context_scope;
 use risingwave_pb::batch_plan::{PbTaskId, PbTaskOutputId, PlanFragment};
 use risingwave_pb::common::BatchQueryEpoch;
+use risingwave_pb::stream_plan::CapturedContext;
 use risingwave_pb::task_service::task_info_response::TaskStatus;
 use risingwave_pb::task_service::{GetDataResponse, TaskInfoResponse};
 use risingwave_pb::PbFieldNotFound;
@@ -425,6 +427,7 @@ impl<C: BatchTaskContext> BatchTaskExecution<C> {
         self: Arc<Self>,
         state_tx: Option<StateReporter>,
         tracing_context: TracingContext,
+        captured_context: CapturedContext,
     ) -> Result<()> {
         let mut state_tx = state_tx;
         trace!(
@@ -433,14 +436,17 @@ impl<C: BatchTaskContext> BatchTaskExecution<C> {
             serde_json::to_string_pretty(self.plan.get_root()?).unwrap()
         );
 
-        let exec = ExecutorBuilder::new(
-            self.plan.root.as_ref().unwrap(),
-            &self.task_id,
-            self.context.clone(),
-            self.epoch.clone(),
-            self.shutdown_rx.clone(),
+        let exec = captured_context_scope!(
+            captured_context.clone(),
+            ExecutorBuilder::new(
+                self.plan.root.as_ref().unwrap(),
+                &self.task_id,
+                self.context.clone(),
+                self.epoch.clone(),
+                self.shutdown_rx.clone(),
+            )
+            .build()
         )
-        .build()
         .await?;
 
         let sender = self.sender.clone();
@@ -472,9 +478,11 @@ impl<C: BatchTaskContext> BatchTaskExecution<C> {
 
                 // We should only pass a reference of sender to execution because we should only
                 // close it after task error has been set.
-                t_1.run(exec, sender, state_tx.as_mut())
-                    .instrument(span)
-                    .await;
+                captured_context_scope!(
+                    captured_context,
+                    t_1.run(exec, sender, state_tx.as_mut()).instrument(span)
+                )
+                .await;
             };
 
             if let Some(batch_metrics) = batch_metrics {

src/batch/src/task/task_manager.rs

@@ -25,6 +25,7 @@ use risingwave_common::util::runtime::BackgroundShutdownRuntime;
 use risingwave_common::util::tracing::TracingContext;
 use risingwave_pb::batch_plan::{PbTaskId, PbTaskOutputId, PlanFragment};
 use risingwave_pb::common::BatchQueryEpoch;
+use risingwave_pb::stream_plan::CapturedContext;
 use risingwave_pb::task_service::task_info_response::TaskStatus;
 use risingwave_pb::task_service::{GetDataResponse, TaskInfoResponse};
 use tokio::sync::mpsc::Sender;
@@ -102,6 +103,7 @@ impl BatchManager {
         context: ComputeNodeContext,
         state_reporter: StateReporter,
         tracing_context: TracingContext,
+        captured_context: CapturedContext,
     ) -> Result<()> {
         trace!("Received task id: {:?}, plan: {:?}", tid, plan);
         let task = BatchTaskExecution::new(tid, plan, context, epoch, self.runtime())?;
@@ -128,7 +130,7 @@ impl BatchManager {
                 task_id,
             );
         };
-        task.async_execute(Some(state_reporter), tracing_context)
+        task.async_execute(Some(state_reporter), tracing_context, captured_context)
             .await
             .inspect_err(|_| {
                 self.cancel_task(&task_id.to_prost());
@@ -151,6 +153,7 @@ impl BatchManager {
             ComputeNodeContext::for_test(),
             StateReporter::new_with_test(),
             TracingContext::none(),
+            CapturedContext::default(),
         )
         .await
     }

src/expr/core/Cargo.toml

@@ -45,6 +45,10 @@ risingwave_udf = { workspace = true }
 smallvec = "1"
 static_assertions = "1"
 thiserror = "1"
+tokio = { version = "0.2", package = "madsim-tokio", features = [
+    "rt",
+    "macros",
+] }
 tracing = "0.1"
 
 [target.'cfg(not(madsim))'.dependencies]

src/expr/core/src/captured_context.rs

@@ -0,0 +1,28 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed 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.
+
+use risingwave_expr::{define_context, Result as ExprResult};
+use risingwave_pb::stream_plan::CapturedContext;
+
+// For all execution mode.
+define_context! {
+    pub TIME_ZONE: String,
+}
+
+pub fn capture_context() -> ExprResult<CapturedContext> {
+    let ctx = TIME_ZONE::try_with(|time_zone| CapturedContext {
+        time_zone: time_zone.to_owned(),
+    })?;
+    Ok(ctx)
+}

src/expr/core/src/lib.rs

@@ -25,6 +25,7 @@
 extern crate self as risingwave_expr;
 
 pub mod aggregate;
+pub mod captured_context;
 #[doc(hidden)]
 pub mod codegen;
 mod error;

src/expr/impl/src/lib.rs

@@ -32,6 +32,7 @@
 #![feature(test)]
 #![feature(arc_unwrap_or_clone)]
 #![feature(iter_array_chunks)]
+#![feature(result_flattening)]
 
 mod aggregate;
 mod scalar;

src/frontend/src/expr/function_impl/make_timestamptz.rs → src/expr/impl/src/scalar/make_timestamptz.rs

@@ -13,19 +13,11 @@
 // limitations under the License.
 
 use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
-use risingwave_common::types::{FloatExt, Timestamptz, F64};
+use risingwave_common::types::{FloatExt, Timestamp, Timestamptz, F64};
+use risingwave_expr::captured_context::TIME_ZONE;
 use risingwave_expr::{capture_context, function, ExprError, Result};
 
-use super::context::TIME_ZONE;
-
-/// Just a wrapper to reuse the `map_err` logic.
-#[inline(always)]
-pub fn time_zone_err(inner_err: String) -> ExprError {
-    ExprError::InvalidParam {
-        name: "time_zone",
-        reason: inner_err.into(),
-    }
-}
+use crate::scalar::timestamptz::timestamp_at_time_zone;
 
 // year int, month int, day int, hour int, min int, sec double precision
 #[function("make_timestamptz(int4, int4, int4, int4, int4, float8) -> timestamptz")]
@@ -64,7 +56,6 @@ fn make_timestamptz_impl(
     min: i32,
     sec: F64,
 ) -> Result<Timestamptz> {
-    let time_zone = Timestamptz::lookup_time_zone(time_zone).map_err(time_zone_err)?;
     if !sec.is_finite() || sec.0.is_sign_negative() {
         return Err(ExprError::InvalidParam {
             name: "sec",
@@ -86,16 +77,6 @@ fn make_timestamptz_impl(
                 reason: "invalid time".into(),
             })?,
     );
-    let date_time = naive_date_time
-        .and_local_timezone(time_zone)
-        .latest()
-        .ok_or_else(|| ExprError::InvalidParam {
-            name: "time_zone",
-            reason: format!(
-                "fail to interpret local timestamp \"{:?}\" in time zone \"{}\"",
-                naive_date_time, time_zone
-            )
-            .into(),
-        })?;
-    Ok(Timestamptz::from(date_time))
+
+    timestamp_at_time_zone(Timestamp(naive_date_time), time_zone)
 }

src/expr/impl/src/scalar/mod.rs

@@ -52,6 +52,7 @@ mod jsonb_info;
 mod jsonb_object;
 mod length;
 mod lower;
+mod make_timestamptz;
 mod md5;
 mod overlay;
 mod position;

src/expr/macro/src/context.rs

@@ -16,7 +16,7 @@ use itertools::Itertools;
 use proc_macro2::TokenStream;
 use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::{Parse, ParseStream};
-use syn::{Error, FnArg, Ident, ItemFn, Result, Token, Type, Visibility};
+use syn::{Error, Expr, FnArg, Ident, ItemFn, Result, Token, Type, Visibility};
 
 use crate::utils::extend_vis_with_super;
 
@@ -224,3 +224,17 @@ pub(super) fn generate_captured_function(
         #new_user_fn
     })
 }
+
+pub(super) struct CapturedContextScopeInput {
+    pub context: Expr,
+    pub closure: Expr,
+}
+
+impl Parse for CapturedContextScopeInput {
+    fn parse(input: ParseStream<'_>) -> Result<Self> {
+        let context: Expr = input.parse()?;
+        input.parse::<Token![,]>()?;
+        let closure: Expr = input.parse()?;
+        Ok(Self { context, closure })
+    }
+}