feat: add DataFile aggregate evaluation (#400)

SuKi2cn · web-flow · commit 428a17146d77 · 2025-12-07T23:12:33.000+08:00
diff --git a/src/iceberg/expression/aggregate.cc b/src/iceberg/expression/aggregate.cc
diff --git a/src/iceberg/expression/aggregate.h b/src/iceberg/expression/aggregate.h
@@ -109,14 +109,15 @@ class ICEBERG_EXPORT BoundAggregate : public Aggregate<BoundTerm>, public Bound
 
     virtual Status Update(const StructLike& data) = 0;
 
-    virtual Status Update(const DataFile& file) {
-      return NotImplemented("Update(DataFile) not implemented");
-    }
+    virtual Status Update(const DataFile& file) = 0;
+
+    /// \brief Whether the aggregator is still valid.
+    virtual bool IsValid() const = 0;
 
     /// \brief Get the result of the aggregation.
     /// \return The result of the aggregation.
     /// \note It is an undefined behavior to call this method if any previous Update call
-    /// has returned an error.
+    /// has returned an error or if IsValid() returns false.
     virtual Literal GetResult() const = 0;
   };
 
@@ -128,6 +129,11 @@ class ICEBERG_EXPORT BoundAggregate : public Aggregate<BoundTerm>, public Bound
 
   Result<Literal> Evaluate(const StructLike& data) const override = 0;
 
+  virtual Result<Literal> Evaluate(const DataFile& file) const = 0;
+
+  /// \brief Whether metrics in the data file are sufficient to evaluate.
+  virtual bool HasValue(const DataFile& file) const = 0;
+
   bool is_bound_aggregate() const override { return true; }
 
   /// \brief Create a new aggregator for this aggregate.
@@ -142,12 +148,15 @@ class ICEBERG_EXPORT BoundAggregate : public Aggregate<BoundTerm>, public Bound
 /// \brief Base class for COUNT aggregates.
 class ICEBERG_EXPORT CountAggregate : public BoundAggregate {
  public:
-  Result<Literal> Evaluate(const StructLike& data) const final;
+  Result<Literal> Evaluate(const StructLike& data) const override;
+  Result<Literal> Evaluate(const DataFile& file) const override;
 
   std::unique_ptr<Aggregator> NewAggregator() const override;
 
   /// \brief Count for a single row. Subclasses implement this.
   virtual Result<int64_t> CountFor(const StructLike& data) const = 0;
+  /// \brief Count using metrics from a data file.
+  virtual Result<int64_t> CountFor(const DataFile& file) const = 0;
 
  protected:
   CountAggregate(Expression::Operation op, std::shared_ptr<BoundTerm> term)
@@ -161,6 +170,8 @@ class ICEBERG_EXPORT CountNonNullAggregate : public CountAggregate {
       std::shared_ptr<BoundTerm> term);
 
   Result<int64_t> CountFor(const StructLike& data) const override;
+  Result<int64_t> CountFor(const DataFile& file) const override;
+  bool HasValue(const DataFile& file) const override;
 
  private:
   explicit CountNonNullAggregate(std::shared_ptr<BoundTerm> term);
@@ -173,6 +184,8 @@ class ICEBERG_EXPORT CountNullAggregate : public CountAggregate {
       std::shared_ptr<BoundTerm> term);
 
   Result<int64_t> CountFor(const StructLike& data) const override;
+  Result<int64_t> CountFor(const DataFile& file) const override;
+  bool HasValue(const DataFile& file) const override;
 
  private:
   explicit CountNullAggregate(std::shared_ptr<BoundTerm> term);
@@ -184,6 +197,8 @@ class ICEBERG_EXPORT CountStarAggregate : public CountAggregate {
   static Result<std::unique_ptr<CountStarAggregate>> Make();
 
   Result<int64_t> CountFor(const StructLike& data) const override;
+  Result<int64_t> CountFor(const DataFile& file) const override;
+  bool HasValue(const DataFile& file) const override;
 
  private:
   CountStarAggregate();
@@ -192,9 +207,11 @@ class ICEBERG_EXPORT CountStarAggregate : public CountAggregate {
 /// \brief Bound MAX aggregate.
 class ICEBERG_EXPORT MaxAggregate : public BoundAggregate {
  public:
-  static std::shared_ptr<MaxAggregate> Make(std::shared_ptr<BoundTerm> term);
+  static Result<std::unique_ptr<MaxAggregate>> Make(std::shared_ptr<BoundTerm> term);
 
   Result<Literal> Evaluate(const StructLike& data) const override;
+  Result<Literal> Evaluate(const DataFile& file) const override;
+  bool HasValue(const DataFile& file) const override;
 
   std::unique_ptr<Aggregator> NewAggregator() const override;
 
@@ -205,9 +222,11 @@ class ICEBERG_EXPORT MaxAggregate : public BoundAggregate {
 /// \brief Bound MIN aggregate.
 class ICEBERG_EXPORT MinAggregate : public BoundAggregate {
  public:
-  static std::shared_ptr<MinAggregate> Make(std::shared_ptr<BoundTerm> term);
+  static Result<std::unique_ptr<MinAggregate>> Make(std::shared_ptr<BoundTerm> term);
 
   Result<Literal> Evaluate(const StructLike& data) const override;
+  Result<Literal> Evaluate(const DataFile& file) const override;
+  bool HasValue(const DataFile& file) const override;
 
   std::unique_ptr<Aggregator> NewAggregator() const override;
 
@@ -234,11 +253,17 @@ class ICEBERG_EXPORT AggregateEvaluator {
   /// \brief Update aggregates with a row.
   virtual Status Update(const StructLike& data) = 0;
 
+  /// \brief Update aggregates using data file metrics.
+  virtual Status Update(const DataFile& file) = 0;
+
   /// \brief Final aggregated value.
   virtual Result<std::span<const Literal>> GetResults() const = 0;
 
   /// \brief Convenience accessor when only one aggregate is evaluated.
   virtual Result<Literal> GetResult() const = 0;
+
+  /// \brief Whether all aggregators are still valid (metrics present).
+  virtual bool AllAggregatorsValid() const = 0;
 };
 
 }  // namespace iceberg
diff --git a/src/iceberg/expression/expressions.cc b/src/iceberg/expression/expressions.cc
@@ -138,6 +138,13 @@ std::shared_ptr<UnboundAggregateImpl<BoundReference>> Expressions::Max(
   return agg;
 }
 
+std::shared_ptr<UnboundAggregateImpl<BoundTransform>> Expressions::Max(
+    std::shared_ptr<UnboundTerm<BoundTransform>> expr) {
+  ICEBERG_ASSIGN_OR_THROW(auto agg, UnboundAggregateImpl<BoundTransform>::Make(
+                                        Expression::Operation::kMax, std::move(expr)));
+  return agg;
+}
+
 std::shared_ptr<UnboundAggregateImpl<BoundReference>> Expressions::Min(std::string name) {
   return Min(Ref(std::move(name)));
 }
@@ -149,6 +156,13 @@ std::shared_ptr<UnboundAggregateImpl<BoundReference>> Expressions::Min(
   return agg;
 }
 
+std::shared_ptr<UnboundAggregateImpl<BoundTransform>> Expressions::Min(
+    std::shared_ptr<UnboundTerm<BoundTransform>> expr) {
+  ICEBERG_ASSIGN_OR_THROW(auto agg, UnboundAggregateImpl<BoundTransform>::Make(
+                                        Expression::Operation::kMin, std::move(expr)));
+  return agg;
+}
+
 // Template implementations for unary predicates
 
 std::shared_ptr<UnboundPredicateImpl<BoundReference>> Expressions::IsNull(
diff --git a/src/iceberg/expression/expressions.h b/src/iceberg/expression/expressions.h
@@ -135,13 +135,19 @@ class ICEBERG_EXPORT Expressions {
   /// \brief Create a MAX aggregate for an unbound term.
   static std::shared_ptr<UnboundAggregateImpl<BoundReference>> Max(
       std::shared_ptr<UnboundTerm<BoundReference>> expr);
+  /// \brief Create a MAX aggregate for an unbound transform term.
+  static std::shared_ptr<UnboundAggregateImpl<BoundTransform>> Max(
+      std::shared_ptr<UnboundTerm<BoundTransform>> expr);
 
   /// \brief Create a MIN aggregate for a field name.
   static std::shared_ptr<UnboundAggregateImpl<BoundReference>> Min(std::string name);
 
   /// \brief Create a MIN aggregate for an unbound term.
   static std::shared_ptr<UnboundAggregateImpl<BoundReference>> Min(
       std::shared_ptr<UnboundTerm<BoundReference>> expr);
+  /// \brief Create a MIN aggregate for an unbound transform term.
+  static std::shared_ptr<UnboundAggregateImpl<BoundTransform>> Min(
+      std::shared_ptr<UnboundTerm<BoundTransform>> expr);
 
   // Unary predicates
 
diff --git a/src/iceberg/expression/term.h b/src/iceberg/expression/term.h
@@ -37,10 +37,13 @@ namespace iceberg {
 /// \brief A term is an expression node that produces a typed value when evaluated.
 class ICEBERG_EXPORT Term : public util::Formattable {
  public:
-  enum class Kind : uint8_t { kReference = 0, kTransform, kExtract };
+  enum class Kind : uint8_t { kReference, kTransform, kExtract };
 
   /// \brief Returns the kind of this term.
   virtual Kind kind() const = 0;
+
+  /// \brief Returns whether this term is unbound.
+  virtual bool is_unbound() const = 0;
 };
 
 template <typename T>
@@ -53,6 +56,8 @@ template <typename B>
 class ICEBERG_EXPORT UnboundTerm : public Unbound<B>, public Term {
  public:
   using BoundType = B;
+
+  bool is_unbound() const override { return true; }
 };
 
 /// \brief Base class for bound terms.
@@ -66,8 +71,6 @@ class ICEBERG_EXPORT BoundTerm : public Bound, public Term {
   /// \brief Returns whether this term may produce null values.
   virtual bool MayProduceNull() const = 0;
 
-  // TODO(gangwu): add a comparator function to Literal and BoundTerm.
-
   /// \brief Returns whether this term is equivalent to another.
   ///
   /// Two terms are equivalent if they produce the same values when evaluated.
@@ -79,6 +82,8 @@ class ICEBERG_EXPORT BoundTerm : public Bound, public Term {
   friend bool operator==(const BoundTerm& lhs, const BoundTerm& rhs) {
     return lhs.Equals(rhs);
   }
+
+  bool is_unbound() const override { return false; }
 };
 
 /// \brief A reference represents a named field in an expression.
diff --git a/src/iceberg/row/struct_like.cc b/src/iceberg/row/struct_like.cc
@@ -19,7 +19,9 @@
 
 #include "iceberg/row/struct_like.h"
 
+#include <string>
 #include <utility>
+#include <vector>
 
 #include "iceberg/result.h"
 #include "iceberg/util/checked_cast.h"
@@ -28,6 +30,44 @@
 
 namespace iceberg {
 
+Result<Scalar> LiteralToScalar(const Literal& literal) {
+  if (literal.IsNull()) {
+    return Scalar{std::monostate{}};
+  }
+
+  switch (literal.type()->type_id()) {
+    case TypeId::kBoolean:
+      return Scalar{std::get<bool>(literal.value())};
+    case TypeId::kInt:
+    case TypeId::kDate:
+      return Scalar{std::get<int32_t>(literal.value())};
+    case TypeId::kLong:
+    case TypeId::kTime:
+    case TypeId::kTimestamp:
+    case TypeId::kTimestampTz:
+      return Scalar{std::get<int64_t>(literal.value())};
+    case TypeId::kFloat:
+      return Scalar{std::get<float>(literal.value())};
+    case TypeId::kDouble:
+      return Scalar{std::get<double>(literal.value())};
+    case TypeId::kString: {
+      const auto& str = std::get<std::string>(literal.value());
+      return Scalar{std::string_view(str)};
+    }
+    case TypeId::kBinary:
+    case TypeId::kFixed: {
+      const auto& bytes = std::get<std::vector<uint8_t>>(literal.value());
+      return Scalar{
+          std::string_view(reinterpret_cast<const char*>(bytes.data()), bytes.size())};
+    }
+    case TypeId::kDecimal:
+      return Scalar{std::get<Decimal>(literal.value())};
+    default:
+      return NotSupported("Cannot convert literal of type {} to Scalar",
+                          literal.type()->ToString());
+  }
+}
+
 StructLikeAccessor::StructLikeAccessor(std::shared_ptr<Type> type,
                                        std::span<const size_t> position_path)
     : type_(std::move(type)) {
diff --git a/src/iceberg/row/struct_like.h b/src/iceberg/row/struct_like.h
@@ -55,6 +55,9 @@ using Scalar = std::variant<std::monostate,  // for null
                             std::shared_ptr<ArrayLike>,   // for list
                             std::shared_ptr<MapLike>>;    // for map
 
+/// \brief Convert a Literal to a Scalar
+Result<Scalar> LiteralToScalar(const Literal& literal);
+
 /// \brief An immutable struct-like wrapper.
 class ICEBERG_EXPORT StructLike {
  public:
diff --git a/src/iceberg/test/aggregate_test.cc b/src/iceberg/test/aggregate_test.cc
