Skip to content

Commit 51a4d66

Browse files
etolbakovetolbakov
authored
feat(to_unixtime): add timestamp types as arguments (#1632)
* feat(to_unixtime): add timestamp types as arguments * feat(to_unixtime): change the return type * feat(to_unixtime): address code review issues * feat(to_unixtime): fix fmt issue
1 parent 1b23815 commit 51a4d66

File tree

3 files changed

+260
-16
lines changed

3 files changed

+260
-16
lines changed

‎src/common/function/src/scalars/timestamp/to_unixtime.rs

Lines changed: 181 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,16 @@ use std::fmt;
1616
use std::str::FromStr;
1717
use std::sync::Arc;
1818

19-
use common_query::error::{self, Result, UnsupportedInputDataTypeSnafu};
19+
use common_query::error::{InvalidFuncArgsSnafu, Result, UnsupportedInputDataTypeSnafu};
2020
use common_query::prelude::{Signature, Volatility};
2121
use common_time::timestamp::TimeUnit;
2222
use common_time::Timestamp;
2323
use datatypes::prelude::ConcreteDataType;
24-
use datatypes::types::StringType;
25-
use datatypes::vectors::{Int64Vector, StringVector, Vector, VectorRef};
24+
use datatypes::types::TimestampType;
25+
use datatypes::vectors::{
26+
Int64Vector, StringVector, TimestampMicrosecondVector, TimestampMillisecondVector,
27+
TimestampNanosecondVector, TimestampSecondVector, Vector, VectorRef,
28+
};
2629
use snafu::ensure;
2730

2831
use crate::scalars::function::{Function, FunctionContext};
@@ -42,26 +45,41 @@ fn convert_to_seconds(arg: &str) -> Option<i64> {
4245
}
4346
}
4447

48+
fn process_vector(vector: &dyn Vector) -> Vec<Option<i64>> {
49+
(0..vector.len())
50+
.map(|i| paste::expr!((vector.get(i)).as_timestamp().map(|ts| ts.value())))
51+
.collect::<Vec<Option<i64>>>()
52+
}
53+
4554
impl Function for ToUnixtimeFunction {
4655
fn name(&self) -> &str {
4756
NAME
4857
}
4958

5059
fn return_type(&self, _input_types: &[ConcreteDataType]) -> Result<ConcreteDataType> {
51-
Ok(ConcreteDataType::timestamp_second_datatype())
60+
Ok(ConcreteDataType::int64_datatype())
5261
}
5362

5463
fn signature(&self) -> Signature {
55-
Signature::exact(
56-
vec![ConcreteDataType::String(StringType)],
64+
Signature::uniform(
65+
1,
66+
vec![
67+
ConcreteDataType::string_datatype(),
68+
ConcreteDataType::int32_datatype(),
69+
ConcreteDataType::int64_datatype(),
70+
ConcreteDataType::timestamp_second_datatype(),
71+
ConcreteDataType::timestamp_millisecond_datatype(),
72+
ConcreteDataType::timestamp_microsecond_datatype(),
73+
ConcreteDataType::timestamp_nanosecond_datatype(),
74+
],
5775
Volatility::Immutable,
5876
)
5977
}
6078

6179
fn eval(&self, _func_ctx: FunctionContext, columns: &[VectorRef]) -> Result<VectorRef> {
6280
ensure!(
6381
columns.len() == 1,
64-
error::InvalidFuncArgsSnafu {
82+
InvalidFuncArgsSnafu {
6583
err_msg: format!(
6684
"The length of the args is not correct, expect exactly one, have: {}",
6785
columns.len()
@@ -79,6 +97,42 @@ impl Function for ToUnixtimeFunction {
7997
.collect::<Vec<_>>(),
8098
)))
8199
}
100+
ConcreteDataType::Int64(_) | ConcreteDataType::Int32(_) => {
101+
let array = columns[0].to_arrow_array();
102+
Ok(Arc::new(Int64Vector::try_from_arrow_array(&array).unwrap()))
103+
}
104+
ConcreteDataType::Timestamp(ts) => {
105+
let array = columns[0].to_arrow_array();
106+
let value = match ts {
107+
TimestampType::Second(_) => {
108+
let vector = paste::expr!(TimestampSecondVector::try_from_arrow_array(
109+
array
110+
)
111+
.unwrap());
112+
process_vector(&vector)
113+
}
114+
TimestampType::Millisecond(_) => {
115+
let vector = paste::expr!(
116+
TimestampMillisecondVector::try_from_arrow_array(array).unwrap()
117+
);
118+
process_vector(&vector)
119+
}
120+
TimestampType::Microsecond(_) => {
121+
let vector = paste::expr!(
122+
TimestampMicrosecondVector::try_from_arrow_array(array).unwrap()
123+
);
124+
process_vector(&vector)
125+
}
126+
TimestampType::Nanosecond(_) => {
127+
let vector = paste::expr!(TimestampNanosecondVector::try_from_arrow_array(
128+
array
129+
)
130+
.unwrap());
131+
process_vector(&vector)
132+
}
133+
};
134+
Ok(Arc::new(Int64Vector::from(value)))
135+
}
82136
_ => UnsupportedInputDataTypeSnafu {
83137
function: NAME,
84138
datatypes: columns.iter().map(|c| c.data_type()).collect::<Vec<_>>(),
@@ -97,28 +151,37 @@ impl fmt::Display for ToUnixtimeFunction {
97151
#[cfg(test)]
98152
mod tests {
99153
use common_query::prelude::TypeSignature;
100-
use datatypes::prelude::ConcreteDataType;
101-
use datatypes::types::StringType;
154+
use datatypes::prelude::{ConcreteDataType, ScalarVectorBuilder};
155+
use datatypes::scalars::ScalarVector;
156+
use datatypes::timestamp::TimestampSecond;
102157
use datatypes::value::Value;
103-
use datatypes::vectors::StringVector;
158+
use datatypes::vectors::{StringVector, TimestampSecondVector};
104159

105160
use super::{ToUnixtimeFunction, *};
106161
use crate::scalars::Function;
107162

108163
#[test]
109-
fn test_to_unixtime() {
164+
fn test_string_to_unixtime() {
110165
let f = ToUnixtimeFunction::default();
111166
assert_eq!("to_unixtime", f.name());
112167
assert_eq!(
113-
ConcreteDataType::timestamp_second_datatype(),
168+
ConcreteDataType::int64_datatype(),
114169
f.return_type(&[]).unwrap()
115170
);
116171

117172
assert!(matches!(f.signature(),
118-
Signature {
119-
type_signature: TypeSignature::Exact(valid_types),
120-
volatility: Volatility::Immutable
121-
} if valid_types == vec![ConcreteDataType::String(StringType)]
173+
Signature {
174+
type_signature: TypeSignature::Uniform(1, valid_types),
175+
volatility: Volatility::Immutable
176+
} if valid_types == vec![
177+
ConcreteDataType::string_datatype(),
178+
ConcreteDataType::int32_datatype(),
179+
ConcreteDataType::int64_datatype(),
180+
ConcreteDataType::timestamp_second_datatype(),
181+
ConcreteDataType::timestamp_millisecond_datatype(),
182+
ConcreteDataType::timestamp_microsecond_datatype(),
183+
ConcreteDataType::timestamp_nanosecond_datatype(),
184+
]
122185
));
123186

124187
let times = vec![
@@ -145,4 +208,106 @@ mod tests {
145208
}
146209
}
147210
}
211+
212+
#[test]
213+
fn test_int_to_unixtime() {
214+
let f = ToUnixtimeFunction::default();
215+
assert_eq!("to_unixtime", f.name());
216+
assert_eq!(
217+
ConcreteDataType::int64_datatype(),
218+
f.return_type(&[]).unwrap()
219+
);
220+
221+
assert!(matches!(f.signature(),
222+
Signature {
223+
type_signature: TypeSignature::Uniform(1, valid_types),
224+
volatility: Volatility::Immutable
225+
} if valid_types == vec![
226+
ConcreteDataType::string_datatype(),
227+
ConcreteDataType::int32_datatype(),
228+
ConcreteDataType::int64_datatype(),
229+
ConcreteDataType::timestamp_second_datatype(),
230+
ConcreteDataType::timestamp_millisecond_datatype(),
231+
ConcreteDataType::timestamp_microsecond_datatype(),
232+
ConcreteDataType::timestamp_nanosecond_datatype(),
233+
]
234+
));
235+
236+
let times = vec![Some(3_i64), None, Some(5_i64), None];
237+
let results = vec![Some(3), None, Some(5), None];
238+
let args: Vec<VectorRef> = vec![Arc::new(Int64Vector::from(times.clone()))];
239+
let vector = f.eval(FunctionContext::default(), &args).unwrap();
240+
assert_eq!(4, vector.len());
241+
for (i, _t) in times.iter().enumerate() {
242+
let v = vector.get(i);
243+
if i == 1 || i == 3 {
244+
assert_eq!(Value::Null, v);
245+
continue;
246+
}
247+
match v {
248+
Value::Int64(ts) => {
249+
assert_eq!(ts, (*results.get(i).unwrap()).unwrap());
250+
}
251+
_ => unreachable!(),
252+
}
253+
}
254+
}
255+
256+
#[test]
257+
fn test_timestamp_to_unixtime() {
258+
let f = ToUnixtimeFunction::default();
259+
assert_eq!("to_unixtime", f.name());
260+
assert_eq!(
261+
ConcreteDataType::int64_datatype(),
262+
f.return_type(&[]).unwrap()
263+
);
264+
265+
assert!(matches!(f.signature(),
266+
Signature {
267+
type_signature: TypeSignature::Uniform(1, valid_types),
268+
volatility: Volatility::Immutable
269+
} if valid_types == vec![
270+
ConcreteDataType::string_datatype(),
271+
ConcreteDataType::int32_datatype(),
272+
ConcreteDataType::int64_datatype(),
273+
ConcreteDataType::timestamp_second_datatype(),
274+
ConcreteDataType::timestamp_millisecond_datatype(),
275+
ConcreteDataType::timestamp_microsecond_datatype(),
276+
ConcreteDataType::timestamp_nanosecond_datatype(),
277+
]
278+
));
279+
280+
let times: Vec<Option<TimestampSecond>> = vec![
281+
Some(TimestampSecond::new(123)),
282+
None,
283+
Some(TimestampSecond::new(42)),
284+
None,
285+
];
286+
let results = vec![Some(123), None, Some(42), None];
287+
let ts_vector: TimestampSecondVector = build_vector_from_slice(&times);
288+
let args: Vec<VectorRef> = vec![Arc::new(ts_vector)];
289+
let vector = f.eval(FunctionContext::default(), &args).unwrap();
290+
assert_eq!(4, vector.len());
291+
for (i, _t) in times.iter().enumerate() {
292+
let v = vector.get(i);
293+
if i == 1 || i == 3 {
294+
assert_eq!(Value::Null, v);
295+
continue;
296+
}
297+
match v {
298+
Value::Int64(ts) => {
299+
assert_eq!(ts, (*results.get(i).unwrap()).unwrap());
300+
}
301+
_ => unreachable!(),
302+
}
303+
}
304+
}
305+
306+
fn build_vector_from_slice<T: ScalarVector>(items: &[Option<T::RefItem<'_>>]) -> T {
307+
let mut builder = T::Builder::with_capacity(items.len());
308+
for item in items {
309+
builder.push(*item);
310+
}
311+
builder.finish()
312+
}
148313
}

‎tests/cases/standalone/common/select/dummy.result

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,3 +42,64 @@ select TO_UNIXTIME('2023-03-01T06:35:02Z');
4242
| 1677652502 |
4343
+-------------------------------------------+
4444

45+
select TO_UNIXTIME(2);
46+
47+
+-----------------------+
48+
| to_unixtime(Int64(2)) |
49+
+-----------------------+
50+
| 2 |
51+
+-----------------------+
52+
53+
create table test_unixtime(a int, b timestamp time index);
54+
55+
Affected Rows: 0
56+
57+
DESC TABLE test_unixtime;
58+
59+
+-------+----------------------+------+---------+---------------+
60+
| Field | Type | Null | Default | Semantic Type |
61+
+-------+----------------------+------+---------+---------------+
62+
| a | Int32 | YES | | FIELD |
63+
| b | TimestampMillisecond | NO | | TIME INDEX |
64+
+-------+----------------------+------+---------+---------------+
65+
66+
insert into test_unixtime values(27, 27);
67+
68+
Affected Rows: 1
69+
70+
select * from test_unixtime;
71+
72+
+----+-------------------------+
73+
| a | b |
74+
+----+-------------------------+
75+
| 27 | 1970-01-01T00:00:00.027 |
76+
+----+-------------------------+
77+
78+
select a from test_unixtime;
79+
80+
+----+
81+
| a |
82+
+----+
83+
| 27 |
84+
+----+
85+
86+
select b from test_unixtime;
87+
88+
+-------------------------+
89+
| b |
90+
+-------------------------+
91+
| 1970-01-01T00:00:00.027 |
92+
+-------------------------+
93+
94+
select TO_UNIXTIME(b) from test_unixtime;
95+
96+
+------------------------------+
97+
| to_unixtime(test_unixtime.b) |
98+
+------------------------------+
99+
| 27 |
100+
+------------------------------+
101+
102+
DROP TABLE test_unixtime;
103+
104+
Affected Rows: 1
105+

‎tests/cases/standalone/common/select/dummy.sql

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,21 @@ select "A";
1111
select * where "a" = "A";
1212

1313
select TO_UNIXTIME('2023-03-01T06:35:02Z');
14+
15+
select TO_UNIXTIME(2);
16+
17+
create table test_unixtime(a int, b timestamp time index);
18+
19+
DESC TABLE test_unixtime;
20+
21+
insert into test_unixtime values(27, 27);
22+
23+
select * from test_unixtime;
24+
25+
select a from test_unixtime;
26+
27+
select b from test_unixtime;
28+
29+
select TO_UNIXTIME(b) from test_unixtime;
30+
31+
DROP TABLE test_unixtime;

0 commit comments

Comments
 (0)