YQL-17471 Add Datetime::EndOf for intervals

commit_hash:98c92942172dfcc73889ff60318242e91828c42b

3 files changed, 164 insertions, 4 deletions

diff --git a/yql/essentials/udfs/common/datetime2/datetime_udf.cpp b/yql/essentials/udfs/common/datetime2/datetime_udf.cpp
index 143eaefbcd..be41a0b8ed 100644
--- a/yql/essentials/udfs/common/datetime2/datetime_udf.cpp
+++ b/yql/essentials/udfs/common/datetime2/datetime_udf.cpp
@@ -1596,7 +1596,25 @@ TValue DoAddYears(const TValue& date, i64 years, const NUdf::IDateBuilder& build
         return storage;
     }
 
-    struct TStartOfBinaryKernelExec : TBinaryKernelExec<TStartOfBinaryKernelExec> {
+    TMaybe<TTMStorage> EndOf(TTMStorage storage, ui64 interval, const IValueBuilder& valueBuilder) {
+        if (interval >= 86400000000ull) {
+            // treat as EndOfDay
+            SetEndOfDay(storage);
+        } else {
+            auto current = storage.ToTimeOfDay();
+            auto rounded = current / interval * (interval + 1) - 1;
+            storage.FromTimeOfDay(rounded);
+        }
+
+        auto& builder = valueBuilder.GetDateBuilder();
+        if (!storage.Validate(builder)) {
+            return {};
+        }
+        return storage;
+    }
+
+    template<bool UseEnd>
+    struct TStartEndOfBinaryKernelExec : TBinaryKernelExec<TStartEndOfBinaryKernelExec<UseEnd>> {
         template<typename TSink>
         static void Process(const IValueBuilder* valueBuilder, TBlockItem arg1, TBlockItem arg2, const TSink& sink) {
             auto& storage = Reference(arg1);
@@ -1606,7 +1624,7 @@ TValue DoAddYears(const TValue& date, i64 years, const NUdf::IDateBuilder& build
                 return;
             }
 
-            if (auto res = StartOf(storage, interval, *valueBuilder)) {
+            if (auto res = (UseEnd ? EndOf : StartOf)(storage, interval, *valueBuilder)) {
                 storage = res.GetRef();
                 sink(arg1);
             } else {
@@ -1627,7 +1645,21 @@ TValue DoAddYears(const TValue& date, i64 years, const NUdf::IDateBuilder& build
         }
         return TUnboxedValuePod{};
     }
-    END_SIMPLE_ARROW_UDF(TStartOf, TStartOfBinaryKernelExec::Do);
+    END_SIMPLE_ARROW_UDF(TStartOf, TStartEndOfBinaryKernelExec<false>::Do);
+
+    BEGIN_SIMPLE_STRICT_ARROW_UDF(TEndOf, TOptional<TResource<TMResourceName>>(TAutoMap<TResource<TMResourceName>>, TAutoMap<TInterval>)) {
+        auto result = args[0];
+        ui64 interval = std::abs(args[1].Get<i64>());
+        if (interval == 0) {
+            return result;
+        }
+        if (auto res = EndOf(Reference(result), interval, *valueBuilder)) {
+            Reference(result) = res.GetRef();
+            return result;
+        }
+        return TUnboxedValuePod{};
+    }
+    END_SIMPLE_ARROW_UDF(TEndOf, TStartEndOfBinaryKernelExec<true>::Do);
 
     struct TTimeOfDayKernelExec : TUnaryKernelExec<TTimeOfDayKernelExec, TReaderTraits::TResource<false>, TFixedSizeArrayBuilder<TDataType<TInterval>::TLayout, false>> {
         template<typename TSink>
diff --git a/yql/essentials/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt b/yql/essentials/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt
index 8dafce2473..5c813a2e6b 100644
--- a/yql/essentials/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt
+++ b/yql/essentials/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt
@@ -352,6 +352,56 @@
                                         "String"
                                     ]
                                 ]
+                            ];
+                            [
+                                "column5";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column6";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column7";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column8";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column9";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
                             ]
                         ]
                     ]
@@ -372,6 +422,19 @@
                         ];
                         [
                             "1970-01-01 23:59:59.999999 Europe/Moscow"
+                        ];
+                        #;
+                        [
+                            "1970-01-01 04:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "1970-01-01 05:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "1970-01-01 05:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "1970-01-01 04:59:57 Europe/Moscow"
                         ]
                     ];
                     [
@@ -389,6 +452,21 @@
                         ];
                         [
                             "2018-12-15 23:59:59.999999 Europe/Moscow"
+                        ];
+                        [
+                            "2018-12-15 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2018-12-15 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2018-12-15 01:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2018-12-15 01:02:00 Europe/Moscow"
+                        ];
+                        [
+                            "2018-12-15 01:01:57 Europe/Moscow"
                         ]
                     ];
                     [
@@ -404,6 +482,21 @@
                         #;
                         [
                             "2105-12-31 23:59:59.999999 GMT"
+                        ];
+                        [
+                            "2105-12-31 13:00:00 GMT"
+                        ];
+                        [
+                            "2105-12-31 16:00:00 GMT"
+                        ];
+                        [
+                            "2105-12-31 16:15:00 GMT"
+                        ];
+                        [
+                            "2105-12-31 16:23:40 GMT"
+                        ];
+                        [
+                            "2105-12-31 16:23:44 GMT"
                         ]
                     ];
                     [
@@ -411,7 +504,22 @@
                         #;
                         #;
                         #;
-                        #
+                        #;
+                        [
+                            "2106-01-01 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2106-01-01 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2106-01-01 01:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2106-01-01 01:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2106-01-01 00:59:58 Europe/Moscow"
+                        ]
                     ];
                     [
                         [
@@ -428,6 +536,21 @@
                         ];
                         [
                             "2019-07-24 23:59:59.999999 Europe/Moscow"
+                        ];
+                        [
+                            "2019-07-24 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2019-07-24 12:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2019-07-24 12:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2019-07-24 12:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2019-07-24 11:59:57 Europe/Moscow"
                         ]
                     ]
                 ]
diff --git a/yql/essentials/udfs/common/datetime2/test/cases/EndOf.sql b/yql/essentials/udfs/common/datetime2/test/cases/EndOf.sql
index ee0eb55e6a..75210b4b95 100644
--- a/yql/essentials/udfs/common/datetime2/test/cases/EndOf.sql
+++ b/yql/essentials/udfs/common/datetime2/test/cases/EndOf.sql
@@ -39,6 +39,11 @@ select
     $format(DateTime::EndOfMonth(`tztimestamp`)),
     $format(DateTime::EndOfWeek(`tztimestamp`)),
     $format(DateTime::EndOfDay(`tztimestamp`)),
+    $format(DateTime::StartOf(`tztimestamp`, Interval("PT13H"))),
+    $format(DateTime::StartOf(`tztimestamp`, Interval("PT4H"))),
+    $format(DateTime::StartOf(`tztimestamp`, Interval("PT15M"))),
+    $format(DateTime::StartOf(`tztimestamp`, Interval("PT20S"))),
+    $format(DateTime::StartOf(`tztimestamp`, Interval("PT7S"))),
 from (
     select
         cast(ftztimestamp as TzTimestamp) as `tztimestamp`