fix: trunc week time zone error

apache · May 17, 2023 · b1b62fb · b1b62fb
1 parent c1b1117
commit b1b62fb
Show file tree

Hide file tree

Showing 3 changed files with 105 additions and 68 deletions.
diff --git a/df_operator/src/udfs/time_bucket.rs b/df_operator/src/udfs/time_bucket.rs
@@ -274,9 +274,14 @@ impl Period {
         granularity: &str,
     ) -> Result<DfColumnarValue> {
         let granularity = DfColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string())));
-        let trunc_array = date_trunc(&[granularity, array.clone()]).context(TruncateTimestamp)?;
+        let array = Self::fix_offset(array)?;
+        let trunc_array = date_trunc(&[granularity, array]).context(TruncateTimestamp)?;
+        let result_array = Self::fix_result_offset(&trunc_array)?;
+        Ok(result_array)
+    }
 
-        let list = match trunc_array {
+    fn fix_result_offset(array: &DfColumnarValue) -> Result<DfColumnarValue> {
+        let list = match array {
             DfColumnarValue::Array(array) => {
                 let array = as_timestamp_nanosecond_array(&array).context(TruncateTimestamp)?;
                 array
@@ -289,6 +294,25 @@ impl Period {
             }
             _ => return UnsupportedScalar.fail(),
         };
+
+        Ok(DfColumnarValue::Array(Arc::new(list)))
+    }
+
+    fn fix_offset(array: &DfColumnarValue) -> Result<DfColumnarValue> {
+        let list = match array {
+            DfColumnarValue::Array(array) => {
+                let array = as_timestamp_nanosecond_array(&array).context(TruncateTimestamp)?;
+                array
+                    .iter()
+                    .map(|ts| {
+                        ts.map(|t| Ok(t + DEFAULT_TIMEZONE_OFFSET_SECS as i64 * 1_000_000_000))
+                            .transpose()
+                    })
+                    .collect::<Result<TimestampNanosecondArray>>()?
+            }
+            _ => return UnsupportedScalar.fail(),
+        };
+
         Ok(DfColumnarValue::Array(Arc::new(list)))
     }
 }
diff --git a/integration_tests/cases/common/function/time_bucket.result b/integration_tests/cases/common/function/time_bucket.result
@@ -20,91 +20,100 @@ VALUES
     (1659577320000, 3),
     (1659571200000, 4),
     (1659484800000, 5),
-    (1656777600000, 6);
+    (1656777600000, 6),
+    (1656898920000, 7);
 
-affected_rows: 6
+affected_rows: 7
 
-SELECT time_bucket(`timestamp`, 'P1Y') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'P1Y') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1Y")),
-Timestamp(1640966400000),
-Timestamp(1640966400000),
-Timestamp(1640966400000),
-Timestamp(1640966400000),
-Timestamp(1640966400000),
-Timestamp(1640966400000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1Y")),
+Timestamp(1656777600000),Timestamp(1640966400000),
+Timestamp(1656898920000),Timestamp(1640966400000),
+Timestamp(1659484800000),Timestamp(1640966400000),
+Timestamp(1659571200000),Timestamp(1640966400000),
+Timestamp(1659577320000),Timestamp(1640966400000),
+Timestamp(1659577422000),Timestamp(1640966400000),
+Timestamp(1659577423000),Timestamp(1640966400000),
 
 
-SELECT time_bucket(`timestamp`, 'P1M') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'P1M') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1M")),
-Timestamp(1656604800000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1M")),
+Timestamp(1656777600000),Timestamp(1656604800000),
+Timestamp(1656898920000),Timestamp(1656604800000),
+Timestamp(1659484800000),Timestamp(1659283200000),
+Timestamp(1659571200000),Timestamp(1659283200000),
+Timestamp(1659577320000),Timestamp(1659283200000),
+Timestamp(1659577422000),Timestamp(1659283200000),
+Timestamp(1659577423000),Timestamp(1659283200000),
 
 
-SELECT time_bucket(`timestamp`, 'P1W') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'P1W') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1W")),
-Timestamp(1656259200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
-Timestamp(1659283200000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1W")),
+Timestamp(1656777600000),Timestamp(1656259200000),
+Timestamp(1656898920000),Timestamp(1656864000000),
+Timestamp(1659484800000),Timestamp(1659283200000),
+Timestamp(1659571200000),Timestamp(1659283200000),
+Timestamp(1659577320000),Timestamp(1659283200000),
+Timestamp(1659577422000),Timestamp(1659283200000),
+Timestamp(1659577423000),Timestamp(1659283200000),
 
 
-SELECT time_bucket(`timestamp`, 'P1D') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'P1D') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1D")),
-Timestamp(1656777600000),
-Timestamp(1659456000000),
-Timestamp(1659542400000),
-Timestamp(1659542400000),
-Timestamp(1659542400000),
-Timestamp(1659542400000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("P1D")),
+Timestamp(1656777600000),Timestamp(1656777600000),
+Timestamp(1656898920000),Timestamp(1656864000000),
+Timestamp(1659484800000),Timestamp(1659456000000),
+Timestamp(1659571200000),Timestamp(1659542400000),
+Timestamp(1659577320000),Timestamp(1659542400000),
+Timestamp(1659577422000),Timestamp(1659542400000),
+Timestamp(1659577423000),Timestamp(1659542400000),
 
 
-SELECT time_bucket(`timestamp`, 'PT1H') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1H') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1H")),
-Timestamp(1656777600000),
-Timestamp(1659484800000),
-Timestamp(1659571200000),
-Timestamp(1659574800000),
-Timestamp(1659574800000),
-Timestamp(1659574800000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1H")),
+Timestamp(1656777600000),Timestamp(1656777600000),
+Timestamp(1656898920000),Timestamp(1656896400000),
+Timestamp(1659484800000),Timestamp(1659484800000),
+Timestamp(1659571200000),Timestamp(1659571200000),
+Timestamp(1659577320000),Timestamp(1659574800000),
+Timestamp(1659577422000),Timestamp(1659574800000),
+Timestamp(1659577423000),Timestamp(1659574800000),
 
 
-SELECT time_bucket(`timestamp`, 'PT1M') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1M') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1M")),
-Timestamp(1656777600000),
-Timestamp(1659484800000),
-Timestamp(1659571200000),
-Timestamp(1659577320000),
-Timestamp(1659577380000),
-Timestamp(1659577380000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1M")),
+Timestamp(1656777600000),Timestamp(1656777600000),
+Timestamp(1656898920000),Timestamp(1656898920000),
+Timestamp(1659484800000),Timestamp(1659484800000),
+Timestamp(1659571200000),Timestamp(1659571200000),
+Timestamp(1659577320000),Timestamp(1659577320000),
+Timestamp(1659577422000),Timestamp(1659577380000),
+Timestamp(1659577423000),Timestamp(1659577380000),
 
 
-SELECT time_bucket(`timestamp`, 'PT1S') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1S') FROM `02_function_time_bucket_table` order by timestamp;
 
-time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1S")),
-Timestamp(1656777600000),
-Timestamp(1659484800000),
-Timestamp(1659571200000),
-Timestamp(1659577320000),
-Timestamp(1659577422000),
-Timestamp(1659577423000),
+timestamp,time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1S")),
+Timestamp(1656777600000),Timestamp(1656777600000),
+Timestamp(1656898920000),Timestamp(1656898920000),
+Timestamp(1659484800000),Timestamp(1659484800000),
+Timestamp(1659571200000),Timestamp(1659571200000),
+Timestamp(1659577320000),Timestamp(1659577320000),
+Timestamp(1659577422000),Timestamp(1659577422000),
+Timestamp(1659577423000),Timestamp(1659577423000),
 
 
 SELECT time_bucket(`timestamp`, 'PT1H', 'yyyy-MM-dd HH:mm:ss') FROM `02_function_time_bucket_table`;
 
 time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1H"),Utf8("yyyy-MM-dd HH:mm:ss")),
 Timestamp(1656777600000),
+Timestamp(1656896400000),
 Timestamp(1659484800000),
 Timestamp(1659571200000),
 Timestamp(1659574800000),
@@ -116,6 +125,7 @@ SELECT time_bucket(`timestamp`, 'PT1H', 'yyyy-MM-dd HH:mm:ss', '+0800') FROM `02
 
 time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1H"),Utf8("yyyy-MM-dd HH:mm:ss"),Utf8("+0800")),
 Timestamp(1656777600000),
+Timestamp(1656896400000),
 Timestamp(1659484800000),
 Timestamp(1659571200000),
 Timestamp(1659574800000),
@@ -127,6 +137,7 @@ SELECT time_bucket(`timestamp`, 'PT1H', 'yyyy-MM-dd HH:mm:ss', '+0800', 'yyyy-MM
 
 time_bucket(02_function_time_bucket_table.timestamp,Utf8("PT1H"),Utf8("yyyy-MM-dd HH:mm:ss"),Utf8("+0800"),Utf8("yyyy-MM-dd HH")),
 Timestamp(1656777600000),
+Timestamp(1656896400000),
 Timestamp(1659484800000),
 Timestamp(1659571200000),
 Timestamp(1659574800000),

diff --git a/integration_tests/cases/common/function/time_bucket.sql b/integration_tests/cases/common/function/time_bucket.sql
@@ -16,16 +16,18 @@ VALUES
     (1659577320000, 3),
     (1659571200000, 4),
     (1659484800000, 5),
-    (1656777600000, 6);
+    (1656777600000, 6),
+    (1656898920000, 7);
+
 
 -- Test all time granularity.
-SELECT time_bucket(`timestamp`, 'P1Y') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'P1M') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'P1W') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'P1D') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'PT1H') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'PT1M') FROM `02_function_time_bucket_table`;
-SELECT time_bucket(`timestamp`, 'PT1S') FROM `02_function_time_bucket_table`;
+SELECT timestamp, time_bucket(`timestamp`, 'P1Y') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'P1M') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'P1W') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'P1D') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1H') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1M') FROM `02_function_time_bucket_table` order by timestamp;
+SELECT timestamp, time_bucket(`timestamp`, 'PT1S') FROM `02_function_time_bucket_table` order by timestamp;
 
 -- Test various parameters.
 -- NOTICE: customizing format is not supported now.