NaT correctness with first aggregator

man-group · Apr 19, 2024 · deb9581 · deb9581
1 parent e7992ff
commit deb9581
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/cpp/arcticdb/processing/aggregation.hpp b/cpp/arcticdb/processing/aggregation.hpp
@@ -398,7 +398,7 @@ class MaxBucketAggregator {
     std::conditional_t<std::is_floating_point_v<T> || TimeType, std::optional<T>,T> max_;
 };
 
-template<typename T>
+template<typename T, bool TimeType = false>
 class FirstBucketAggregator {
 public:
     void push(T value) {
@@ -410,6 +410,10 @@ class FirstBucketAggregator {
             if (ARCTICDB_UNLIKELY(!has_value_ || std::isnan(first_))) {
                 first_ = value;
             }
+        } else if constexpr (std::is_same_v<T, timestamp> && TimeType) {
+            if (ARCTICDB_UNLIKELY(!has_value_ || first_ == NaT)) {
+                first_ = value;
+            }
         } else {
             if (ARCTICDB_UNLIKELY(!has_value_)) {
                 first_ = value;
@@ -426,6 +430,8 @@ class FirstBucketAggregator {
         } else {
             if constexpr (std::is_floating_point_v<T>) {
                 res = std::numeric_limits<T>::quiet_NaN();
+            } else if constexpr(std::is_same_v<T, timestamp> && TimeType) {
+                res = NaT;
             } else {
                 internal::raise<ErrorCode::E_ASSERTION_FAILURE>("FirstBucketAggregator::finalize called with no values pushed");
             }
@@ -565,7 +571,11 @@ class SortedAggregator
             }
         } else if constexpr (aggregation_operator == SortedAggregationOperator::FIRST) {
             if constexpr (is_numeric_type(scalar_type_info::data_type) || is_bool_type(scalar_type_info::data_type)) {
-                return FirstBucketAggregator<typename scalar_type_info::RawType>();
+                if constexpr (is_time_type(scalar_type_info::data_type)) {
+                    return FirstBucketAggregator<typename scalar_type_info::RawType, true>();
+                } else {
+                    return FirstBucketAggregator<typename scalar_type_info::RawType>();
+                }
             } else if (is_sequence_type(scalar_type_info::data_type)) {
                 return FirstBucketAggregator<std::optional<std::string_view>>();
             }

diff --git a/python/tests/unit/arcticdb/version_store/test_resample.py b/python/tests/unit/arcticdb/version_store/test_resample.py
@@ -334,9 +334,9 @@ def test_resampling_datetime_columns(lmdb_version_store_v1):
     expected = df.resample("us").agg(
         # mean=pd.NamedAgg("col", "mean"),
         # min=pd.NamedAgg("col", "min"),
-        max=pd.NamedAgg("col", "max"),
+        # max=pd.NamedAgg("col", "max"),
         # first=pd.NamedAgg("col", "first"),
-        # last=pd.NamedAgg("col", "last"),
+        last=pd.NamedAgg("col", "last"),
         # count=pd.NamedAgg("col", "count"),
     )
     expected = expected.reindex(columns=sorted(expected.columns))
@@ -347,9 +347,9 @@ def test_resampling_datetime_columns(lmdb_version_store_v1):
         {
             # "mean": ("col", "mean"),
             # "min": ("col", "min"),
-            "max": ("col", "max"),
+            # "max": ("col", "max"),
             # "first": ("col", "first"),
-            # "last": ("col", "last"),
+            "last": ("col", "last"),
             # "count": ("col", "count"),
         }
     )