starred/openobserve
1
0
Fork 0
mirror of https://github.com/openobserve/openobserve.git synced 2026-11-06 16:09:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement cast_to_timestamp (#6084)

Browse Source
This commit is contained in:
Hengfei Yang
2025-02-25 00:40:08 +08:00
committed by GitHub
parent 50daaec8dc
commit bfbb4fe1fa
3 changed files with 223 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/service/search/datafusion/exec.rs
+1 -1
View File
@@ -468,7 +468,7 @@ pub fn register_udf(ctx: &SessionContext, org_id: &str) -> Result<()> {
ctx.register_udaf(AggregateUDF::from(
super::udaf::percentile_cont::PercentileCont::new(),
));
ctx.register_udf(super::udf::cast_to_timestamp_udf::CAST_TO_TIMESTAMP_UDF.clone());
let udf_list = get_all_transform(org_id)?;
for udf in udf_list {
ctx.register_udf(udf.clone());
src/service/search/datafusion/udf/cast_to_timestamp_udf.rs
+215
View File
@@ -0,0 +1,215 @@
// Copyright 2025 OpenObserve Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
use std::any::Any;
use arrow_schema::TimeUnit;
use datafusion::{
arrow::datatypes::DataType,
common::{exec_err, not_impl_err},
error::Result,
functions::datetime::to_timestamp::ToTimestampMicrosFunc,
logical_expr::{ColumnarValue, ScalarUDF, ScalarUDFImpl, Signature, Volatility},
};
use once_cell::sync::Lazy;
/// The name of the cast_to_timestamp UDF given to DataFusion.
pub const CAST_TO_TIMESTAMP_UDF_NAME: &str = "cast_to_timestamp";
/// Implementation of cast_to_timestamp
pub(crate) static CAST_TO_TIMESTAMP_UDF: Lazy<ScalarUDF> =
Lazy::new(|| ScalarUDF::from(CastToTimestampUdf::new()));
#[derive(Debug, Clone)]
struct CastToTimestampUdf {
signature: Signature,
}
impl CastToTimestampUdf {
fn new() -> Self {
Self {
signature: Signature::variadic_any(Volatility::Immutable),
}
}
}
impl ScalarUDFImpl for CastToTimestampUdf {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
CAST_TO_TIMESTAMP_UDF_NAME
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Int64)
}
fn invoke(&self, _args: &[ColumnarValue]) -> Result<ColumnarValue> {
not_impl_err!(
"Function {} does not implement invoke but called",
self.name()
)
}
fn invoke_batch(&self, args: &[ColumnarValue], number_rows: usize) -> Result<ColumnarValue> {
let ttm = ToTimestampMicrosFunc::new();
let ret = ttm.invoke_batch(args, number_rows)?;
match ret.data_type() {
DataType::Int32 | DataType::Int64 | DataType::Null | DataType::Float64 => {
ret.cast_to(&DataType::Int64, None)
}
DataType::Date64 | DataType::Date32 | DataType::Timestamp(_, None) => ret
.cast_to(&DataType::Timestamp(TimeUnit::Microsecond, None), None)?
.cast_to(&DataType::Int64, None),
DataType::Timestamp(_, Some(tz)) => ret
.cast_to(&DataType::Timestamp(TimeUnit::Microsecond, Some(tz)), None)?
.cast_to(&DataType::Int64, None),
DataType::Utf8 => CastToTimestampUdf::new()
.invoke_batch(args, number_rows)?
.cast_to(&DataType::Int64, None),
other => {
exec_err!(
"Unsupported data type {:?} for function {}",
other,
self.name()
)
}
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use arrow::array::{Int64Array, StringArray};
use datafusion::{
arrow::{
datatypes::{Field, Schema},
record_batch::RecordBatch,
},
assert_batches_eq,
datasource::MemTable,
prelude::SessionContext,
};
use super::*;
#[tokio::test]
async fn test_cast_to_timestamp_udf() {
let now = config::utils::time::now_micros();
let sqls = [
(
"SELECT 1, LENGTH(cast_to_timestamp(NOW())) AS len from t WHERE time <= cast_to_timestamp(NOW())",
vec![
"+----------+-----+",
"| Int64(1) | len |",
"+----------+-----+",
"| 1 | 16 |",
"+----------+-----+",
],
),
(
"SELECT 1 from t WHERE time > cast_to_timestamp(NOW() - INTERVAL '1 minute')",
vec![
"+----------+",
"| Int64(1) |",
"+----------+",
"| 1 |",
"+----------+",
],
),
(
"SELECT 1 from t WHERE time > cast_to_timestamp(NOW() - INTERVAL '1 hour')",
vec![
"+----------+",
"| Int64(1) |",
"+----------+",
"| 1 |",
"+----------+",
],
),
(
"SELECT 1 from t WHERE time > cast_to_timestamp('2025-01-01T00:00:00.000Z')",
vec![
"+----------+",
"| Int64(1) |",
"+----------+",
"| 1 |",
"+----------+",
],
),
(
"SELECT 1 from t WHERE time > cast_to_timestamp('2025-01-01T00:00:00+08:00')",
vec![
"+----------+",
"| Int64(1) |",
"+----------+",
"| 1 |",
"+----------+",
],
),
(
"SELECT 1 from t WHERE time > cast_to_timestamp(1740398708294000)",
vec![
"+----------+",
"| Int64(1) |",
"+----------+",
"| 1 |",
"+----------+",
],
),
];
// define a schema.
let schema = Arc::new(Schema::new(vec![
Field::new("log", DataType::Utf8, false),
Field::new("time", DataType::Int64, false),
]));
// define data.
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(StringArray::from(vec!["a"])),
Arc::new(Int64Array::from(vec![now])),
],
)
.unwrap();
// declare a new context. In spark API, this corresponds to a new spark
// SQLsession
let ctx = SessionContext::new();
ctx.register_udf(CAST_TO_TIMESTAMP_UDF.clone());
// declare a table in memory. In spark API, this corresponds to
// createDataFrame(...).
let provider = MemTable::try_new(schema, vec![vec![batch]]).unwrap();
ctx.register_table("t", Arc::new(provider)).unwrap();
for item in sqls {
let df = ctx.sql(item.0).await.unwrap();
let data = df.collect().await.unwrap();
assert_batches_eq!(item.1, &data);
}
}
}
src/service/search/datafusion/udf/mod.rs
+7 -2
View File
@@ -1,4 +1,4 @@
// Copyright 2024 OpenObserve Inc.
// Copyright 2025 OpenObserve Inc.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
@@ -22,6 +22,7 @@ pub(crate) mod arrjoin_udf;
pub(crate) mod arrsort_udf;
pub(crate) mod arrzip_udf;
pub(crate) mod cast_to_arr_udf;
pub(crate) mod cast_to_timestamp_udf;
#[cfg(feature = "enterprise")]
pub(crate) mod cipher_udf;
pub(crate) mod date_format_udf;
@@ -50,7 +51,7 @@ pub(crate) const REGEX_NOT_MATCH_UDF_NAME: &str = "re_not_match";
/// The name of the regex_matches UDF given to DataFusion.
pub(crate) const REGEX_MATCHES_UDF_NAME: &str = "re_matches";
pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 10] = [
pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 11] = [
ZoFunction {
name: "match_all_raw",
text: "match_all_raw('v')",
@@ -91,6 +92,10 @@ pub(crate) const DEFAULT_FUNCTIONS: [ZoFunction; 10] = [
name: REGEX_MATCHES_UDF_NAME,
text: "re_matches(field, 'pattern')",
},
ZoFunction {
name: cast_to_timestamp_udf::CAST_TO_TIMESTAMP_UDF_NAME,
text: "cast_to_timestamp('pattern')",
},
];
pub fn stringify_json_value(field: &json::Value) -> String {