rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-05-31 20:30:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: make query be aware of timezone setting (#3175)

Browse Source 
* feat: let TypeConversionRule aware query context timezone setting

* chore: don't optimize explain command

* feat: parse string into timestamp with timezone

* fix: compile error

* chore: check the scalar value type in predicate

* chore: remove mut for engine context

* chore: return none if the scalar value is utf8 in time range predicate

* fix: some fixme

* feat: let Date and DateTime parsing from string value be aware of timezone

* chore: tweak

* test: add datetime from_str test with timezone

* feat: construct function context from query context

* test: add timezone test for to_unixtime and date_format function

* fix: typo

* chore: apply suggestion

* test: adds string with timezone

* chore: apply CR suggestion

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>

* chore: apply suggestion

---------

Co-authored-by: Lei, HUANG <6406592+v0y4g3r@users.noreply.github.com>
This commit is contained in:
dennis zhuang
2024-01-22 22:14:03 +08:00
committed by GitHub
parent 2bf4b08a6b
commit 6a12c27e78
45 changed files with 995 additions and 238 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/common/function/Cargo.toml
View File

@@ -19,6 +19,7 @@ num = "0.4"
num-traits = "0.2"
once_cell.workspace = true
paste = "1.0"
session.workspace = true
snafu.workspace = true
statrs = "0.16"

8
src/common/function/src/function.rs
View File

@@ -17,20 +17,20 @@ use std::sync::Arc;
use common_query::error::Result;
use common_query::prelude::Signature;
use common_time::timezone::get_timezone;
use common_time::Timezone;
use datatypes::data_type::ConcreteDataType;
use datatypes::vectors::VectorRef;
use session::context::{QueryContextBuilder, QueryContextRef};
/// The function execution context
#[derive(Clone)]
pub struct FunctionContext {
pub timezone: Timezone,
pub query_ctx: QueryContextRef,
}
impl Default for FunctionContext {
fn default() -> Self {
Self {
timezone: get_timezone(None).clone(),
query_ctx: QueryContextBuilder::default().build(),
}
}
}

6
src/common/function/src/scalars/date/date_format.rs
View File

@@ -79,7 +79,7 @@ impl Function for DateFormatFunction {
let result = match (ts, format) {
(Some(ts), Some(fmt)) => Some(
ts.as_formatted_string(&fmt, Some(&func_ctx.timezone))
ts.as_formatted_string(&fmt, Some(&func_ctx.query_ctx.timezone()))
.map_err(BoxedError::new)
.context(error::ExecuteSnafu)?,
),
@@ -96,7 +96,7 @@ impl Function for DateFormatFunction {
let result = match (date, format) {
(Some(date), Some(fmt)) => date
.as_formatted_string(&fmt, Some(&func_ctx.timezone))
.as_formatted_string(&fmt, Some(&func_ctx.query_ctx.timezone()))
.map_err(BoxedError::new)
.context(error::ExecuteSnafu)?,
_ => None,
@@ -112,7 +112,7 @@ impl Function for DateFormatFunction {
let result = match (datetime, format) {
(Some(datetime), Some(fmt)) => datetime
.as_formatted_string(&fmt, Some(&func_ctx.timezone))
.as_formatted_string(&fmt, Some(&func_ctx.query_ctx.timezone()))
.map_err(BoxedError::new)
.context(error::ExecuteSnafu)?,
_ => None,

9
src/common/function/src/scalars/timestamp/greatest.rs
View File

@@ -104,7 +104,6 @@ impl fmt::Display for GreatestFunction {
#[cfg(test)]
mod tests {
use std::str::FromStr;
use std::sync::Arc;
use common_time::Date;
@@ -137,11 +136,11 @@ mod tests {
assert_eq!(result.len(), 2);
assert_eq!(
result.get(0),
Value::Date(Date::from_str("2001-02-01").unwrap())
Value::Date(Date::from_str_utc("2001-02-01").unwrap())
);
assert_eq!(
result.get(1),
Value::Date(Date::from_str("2012-12-23").unwrap())
Value::Date(Date::from_str_utc("2012-12-23").unwrap())
);
}
@@ -162,11 +161,11 @@ mod tests {
assert_eq!(result.len(), 2);
assert_eq!(
result.get(0),
Value::Date(Date::from_str("1970-01-01").unwrap())
Value::Date(Date::from_str_utc("1970-01-01").unwrap())
);
assert_eq!(
result.get(1),
Value::Date(Date::from_str("1970-01-03").unwrap())
Value::Date(Date::from_str_utc("1970-01-03").unwrap())
);
}
}

14
src/common/function/src/scalars/timestamp/to_unixtime.rs
View File

@@ -13,7 +13,6 @@
// limitations under the License.
use std::fmt;
use std::str::FromStr;
use std::sync::Arc;
use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
@@ -31,16 +30,17 @@ pub struct ToUnixtimeFunction;
const NAME: &str = "to_unixtime";
fn convert_to_seconds(arg: &str) -> Option<i64> {
if let Ok(dt) = DateTime::from_str(arg) {
fn convert_to_seconds(arg: &str, func_ctx: &FunctionContext) -> Option<i64> {
let timezone = &func_ctx.query_ctx.timezone();
if let Ok(dt) = DateTime::from_str(arg, Some(timezone)) {
return Some(dt.val() / 1000);
}
if let Ok(ts) = Timestamp::from_str(arg) {
if let Ok(ts) = Timestamp::from_str(arg, Some(timezone)) {
return Some(ts.split().0);
}
if let Ok(date) = Date::from_str(arg) {
if let Ok(date) = Date::from_str(arg, Some(timezone)) {
return Some(date.to_secs());
}
@@ -92,7 +92,7 @@ impl Function for ToUnixtimeFunction {
)
}
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
fn eval(&self, func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
ensure!(
columns.len() == 1,
InvalidFuncArgsSnafu {
@@ -108,7 +108,7 @@ impl Function for ToUnixtimeFunction {
match columns[0].data_type() {
ConcreteDataType::String(_) => Ok(Arc::new(Int64Vector::from(
(0..vector.len())
.map(|i| convert_to_seconds(&vector.get(i).to_string()))
.map(|i| convert_to_seconds(&vector.get(i).to_string(), &func_ctx))
.collect::<Vec<_>>(),
))),
ConcreteDataType::Int64(_) | ConcreteDataType::Int32(_) => {

16
src/common/function/src/scalars/udf.rs
View File

@@ -21,22 +21,24 @@ use common_query::prelude::{
use datatypes::error::Error as DataTypeError;
use datatypes::prelude::*;
use datatypes::vectors::Helper;
use session::context::QueryContextRef;
use snafu::ResultExt;
use crate::function::{FunctionContext, FunctionRef};
/// Create a ScalarUdf from function.
pub fn create_udf(func: FunctionRef) -> ScalarUdf {
/// Create a ScalarUdf from function and query context.
pub fn create_udf(func: FunctionRef, query_ctx: QueryContextRef) -> ScalarUdf {
let func_cloned = func.clone();
let return_type: ReturnTypeFunction = Arc::new(move |input_types: &[ConcreteDataType]| {
Ok(Arc::new(func_cloned.return_type(input_types)?))
});
let func_cloned = func.clone();
let fun: ScalarFunctionImplementation = Arc::new(move |args: &[ColumnarValue]| {
// FIXME(dennis): set the timezone into function context
// Question: how to get the timezone from the query context?
let func_ctx = FunctionContext::default();
let func_ctx = FunctionContext {
query_ctx: query_ctx.clone(),
};
let len = args
.iter()
@@ -72,6 +74,7 @@ mod tests {
use datatypes::prelude::{ScalarVector, Vector, VectorRef};
use datatypes::value::Value;
use datatypes::vectors::{BooleanVector, ConstantVector};
use session::context::QueryContextBuilder;
use super::*;
use crate::function::Function;
@@ -80,6 +83,7 @@ mod tests {
#[test]
fn test_create_udf() {
let f = Arc::new(TestAndFunction);
let query_ctx = QueryContextBuilder::default().build();
let args: Vec<VectorRef> = vec![
Arc::new(ConstantVector::new(
@@ -97,7 +101,7 @@ mod tests {
}
// create a udf and test it again
let udf = create_udf(f.clone());
let udf = create_udf(f.clone(), query_ctx);
assert_eq!("test_and", udf.name);
assert_eq!(f.signature(), udf.signature);