Add UDF DateTime::EndOfMonth

4 files changed, 355 insertions, 3 deletions

diff --git a/ydb/library/yql/udfs/common/datetime2/datetime_udf.cpp b/ydb/library/yql/udfs/common/datetime2/datetime_udf.cpp
index 408ab4f6e9f..59e66aab839 100644
--- a/ydb/library/yql/udfs/common/datetime2/datetime_udf.cpp
+++ b/ydb/library/yql/udfs/common/datetime2/datetime_udf.cpp
@@ -171,7 +171,7 @@ public:
     }
 };
 
-template <const char* TFuncName, typename TFieldStorage, TFieldStorage (*FieldFunc)(const TUnboxedValuePod&), ui32 Divisor, ui32 Scale, ui32 Limit, bool Fractional> 
+template <const char* TFuncName, typename TFieldStorage, TFieldStorage (*FieldFunc)(const TUnboxedValuePod&), ui32 Divisor, ui32 Scale, ui32 Limit, bool Fractional>
 struct TGetTimeComponent {
     typedef bool TTypeAwareMarker;
 
@@ -261,8 +261,8 @@ struct TGetTimeComponent {
                     }
 
                     auto typeId = data.GetTypeId();
-                    if (typeId == TDataType<TDate>::Id || 
-                        typeId == TDataType<TDatetime>::Id || 
+                    if (typeId == TDataType<TDate>::Id ||
+                        typeId == TDataType<TDatetime>::Id ||
                         typeId == TDataType<TTimestamp>::Id) {
 
                         builder.Args()->Add(argsTuple.GetElementType(0)).Done();
@@ -995,6 +995,22 @@ NUdf::TUnboxedValuePod DoAddYears(const NUdf::TUnboxedValuePod& date, i64 years,
         return result;
     }
 
+    SIMPLE_STRICT_UDF(TEndOfMonth, TOptional<TResource<TMResourceName>>(TAutoMap<TResource<TMResourceName>>)) {
+        auto result = args[0];
+        auto& storage = Reference(result);
+        storage.Day = NMiniKQL::GetMonthLength(storage.Month, NMiniKQL::IsLeapYear(storage.Year));
+        storage.Hour = 0;
+        storage.Minute = 0;
+        storage.Second = 0;
+        storage.Microsecond = 0;
+
+        auto& builder = valueBuilder->GetDateBuilder();
+        if (!storage.Validate(builder)) {
+            return TUnboxedValuePod();
+        }
+        return result;
+    }
+
     SIMPLE_STRICT_UDF(TStartOfWeek, TOptional<TResource<TMResourceName>>(TAutoMap<TResource<TMResourceName>>)) {
         auto result = args[0];
         auto& storage = Reference(result);
@@ -1787,6 +1803,8 @@ NUdf::TUnboxedValuePod DoAddYears(const NUdf::TUnboxedValuePod& date, i64 years,
         TShiftQuarters,
         TShiftMonths,
 
+        TEndOfMonth,
+
         TToUnits<ToSecondsName, ui32, 1>,
         TToUnits<ToMillisecondsName, ui64, 1000>,
         TToUnits<ToMicrosecondsName, ui64, 1000000>,
diff --git a/ydb/library/yql/udfs/common/datetime2/test/canondata/result.json b/ydb/library/yql/udfs/common/datetime2/test/canondata/result.json
index 4291f7e1fa1..9c70db5e3bd 100644
--- a/ydb/library/yql/udfs/common/datetime2/test/canondata/result.json
+++ b/ydb/library/yql/udfs/common/datetime2/test/canondata/result.json
@@ -9,6 +9,11 @@
             "uri": "file://test.test_BlockTo_/results.txt"
         }
     ],
+    "test.test[EndOf]": [
+        {
+            "uri": "file://test.test_EndOf_/results.txt"
+        }
+    ],
     "test.test[Format]": [
         {
             "uri": "file://test.test_Format_/results.txt"
diff --git a/ydb/library/yql/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt b/ydb/library/yql/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt
new file mode 100644
index 00000000000..508cd4438ed
--- /dev/null
+++ b/ydb/library/yql/udfs/common/datetime2/test/canondata/test.test_EndOf_/results.txt
@@ -0,0 +1,295 @@
+[
+    {
+        "Label" = "Normal cases";
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column2";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column3";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column4";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "2023-07-31 00:00:00 Europe/Moscow"
+                        ];
+                        [
+                            "2023-08-31 00:00:00 GMT"
+                        ];
+                        [
+                            "2023-09-30 00:00:00 GMT"
+                        ];
+                        [
+                            "2023-02-28 00:00:00 GMT"
+                        ];
+                        [
+                            "2024-02-29 00:00:00 GMT"
+                        ]
+                    ]
+                ]
+            }
+        ]
+    };
+    {
+        "Label" = "Minimal timestamp value";
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column2";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "1970-01-01 00:00:00 GMT"
+                        ];
+                        [
+                            "1970-01-31 00:00:00 GMT"
+                        ];
+                        #
+                    ]
+                ]
+            }
+        ]
+    };
+    {
+        "Label" = "Maximum timestamp value";
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column2";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column3";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "2105-12-31 23:59:59.999999 GMT"
+                        ];
+                        [
+                            "2105-12-31 00:00:00 GMT"
+                        ];
+                        [
+                            "2105-12-31 00:00:00 GMT"
+                        ];
+                        #
+                    ]
+                ]
+            }
+        ]
+    };
+    {
+        "Label" = "Timestamp below minimum";
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        [
+                            "1969-12-31 23:59:59.999999 Atlantic/Azores"
+                        ];
+                        #
+                    ]
+                ]
+            }
+        ]
+    };
+    {
+        "Label" = "Timestamp above maximum";
+        "Write" = [
+            {
+                "Type" = [
+                    "ListType";
+                    [
+                        "StructType";
+                        [
+                            [
+                                "column0";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ];
+                            [
+                                "column1";
+                                [
+                                    "OptionalType";
+                                    [
+                                        "DataType";
+                                        "String"
+                                    ]
+                                ]
+                            ]
+                        ]
+                    ]
+                ];
+                "Data" = [
+                    [
+                        #;
+                        #
+                    ]
+                ]
+            }
+        ]
+    }
+]
\ No newline at end of file
diff --git a/ydb/library/yql/udfs/common/datetime2/test/cases/EndOf.sql b/ydb/library/yql/udfs/common/datetime2/test/cases/EndOf.sql
new file mode 100644
index 00000000000..61b4a29e536
--- /dev/null
+++ b/ydb/library/yql/udfs/common/datetime2/test/cases/EndOf.sql
@@ -0,0 +1,34 @@
+/* syntax version 1 */
+$format = DateTime::Format("%Y-%m-%d %H:%M:%S %Z");
+
+select
+    $format(DateTime::EndOfMonth(TzDateTime('2023-07-07T01:02:03,Europe/Moscow'))),
+    $format(DateTime::EndOfMonth(Date('2023-08-08'))),
+    $format(DateTime::EndOfMonth(Date('2023-09-09'))),
+    $format(DateTime::EndOfMonth(Date('2023-02-02'))),
+    $format(DateTime::EndOfMonth(Date('2024-02-02')))
+into result `Normal cases`;
+
+$tsMin = '1970-01-01T00:00:00.000000';
+$tsMax = '2105-12-31T23:59:59.999999';
+$tsBelow = '1969-12-31T23:59:59.999999';
+$tsAbove = '2106-01-01T00:00:00.000000';
+
+select $format(cast($tsMin || 'Z' as Timestamp))
+    , $format(DateTime::EndOfMonth(cast($tsMin || 'Z' as Timestamp)))
+    , $format(DateTime::EndOfMonth(cast($tsMin || ',Atlantic/Madeira' as Timestamp)))
+into result `Minimal timestamp value`;
+
+select $format(cast($tsMax || 'Z' as Timestamp))
+    , $format(DateTime::EndOfMonth(cast($tsMax || 'Z' as Timestamp)))
+    , $format(DateTime::EndOfMonth(cast('2105-12-12T00:00:00Z' as Timestamp)))
+    , $format(DateTime::EndOfMonth(cast($tsMax || ',Atlantic/Azores' as Timestamp)))
+into result `Maximum timestamp value`;
+
+select $format(cast($tsBelow || ',Atlantic/Azores' as TzTimestamp))
+    , $format(DateTime::EndOfMonth(cast($tsBelow || ',Atlantic/Azores' as TzTimestamp)))
+into result `Timestamp below minimum`;
+
+select $format(cast($tsAbove || ',Atlantic/Madeira' as TzTimestamp))
+    , $format(DateTime::EndOfMonth(cast($tsAbove || ',Atlantic/Madeira' as TzTimestamp)))
+into result `Timestamp above maximum`;