Track Value constraints through plan construction

Summary:
This diff implements end-to-end constraint propagation through the Axiom query plan construction process. Building on the filter selectivity estimation framework (D88164653), this change tracks refined Value constraints (min/max ranges, cardinalities, null fractions) from leaf scans through joins, filters, projects, and aggregations, enabling more accurate cardinality estimates at every level of the plan.

The core motivation is to improve cost-based optimization decisions by maintaining up-to-date Value metadata as constraints are derived from predicates and joins. For example, after a filter `WHERE age > 18 AND age < 65`, the optimizer now knows the age column has min=18, max=65, and reduced cardinality, which influences downstream join cost estimates. Similarly, join equalities like `customer.id = order.customer_id` update both columns' constraints to reflect the intersection of their ranges and cardinalities.

This constraint tracking is essential for:
- Accurate join cardinality estimation (especially for multi-way joins)
- Proper null handling in outer joins (tracking which columns become nullable)
- Aggregate cardinality estimation (result size of GROUP BY)
- Project cardinality propagation (computing distinct counts through expressions)
- Filter selectivity estimation (more accurate with refined input constraints)

Main changes:

1. **PlanState constraint tracking** (Plan.h, Plan.cpp):
   - Added `ConstraintMap constraints` field to PlanState to track Value constraints by expression ID during plan construction
   - Added `exprConstraint()` function that derives constraints for expressions based on their type:
     - Literals: use the literal's min/max (already set in expr->value())
     - Columns: look up in state.constraints, fall back to expr->value()
     - Field access: propagate cardinality from base expression
     - Aggregates: set cardinality from state.cost.cardinality (group count)
     - Function calls: use functionConstraint metadata if available, otherwise max of argument cardinalities
   - Modified Plan constructor to populate Plan::constraints from state.constraints using QueryGraphContext::registerAny() for lifetime management
   - Modified NextJoin to preserve constraints across join candidates
   - Added PlanStateSaver to save/restore constraints along with cost and placed sets

2. **Join constraint propagation** (RelationOp.cpp):
   - Added `addJoinConstraint()` function to update constraints for join key pairs:
     - For inner joins: both sides become non-nullable (nullFraction=0), ranges intersect via columnComparisonSelectivity
     - For outer joins: optional side becomes nullable, nullFraction set to (1 - innerFanout)
     - Non-key columns from optional side get nullFraction = (1 - innerFanout)
   - Added `addJoinConstraints()` to apply constraint updates for all key pairs
   - Modified Join constructor to:
     - Accept new `innerFanout` parameter (fanout if this were an inner join, used for null fraction calculation)
     - Call addJoinConstraints() with join type information
     - Handle filter expressions using conjunctsSelectivity() to get filter selectivity
     - For semi-joins and anti-joins, multiply fanout by filter selectivity
     - For semi-project (mark joins), update the mark column's trueFraction
     - Propagate non-key nullable constraints for outer joins

3. **Filter constraint integration** (Filters.cpp, Filters.h):
   - Modified `value()` function to first check state.constraints before falling back to expr->value()
   - Added `addConstraint()` helper that enforces type-based cardinality limits:
     - BOOLEAN: max 2 distinct values
     - TINYINT: max 256 distinct values
     - SMALLINT: max 65,536 distinct values
   - Modified `conjunctsSelectivity()` to call `exprConstraint()` for each conjunct before computing selectivity, ensuring constraints are computed and cached
   - Changed `exprSelectivity()` and `conjunctsSelectivity()` signatures to take non-const `PlanState&` to allow updating state.constraints
   - Added `constraintsString()` debugging helper to format ConstraintMap as readable string

4. **Project constraint propagation** (RelationOp.cpp):
   - Modified Project constructor to accept PlanState parameter
   - Added loop to derive and store constraints for each output column by calling exprConstraint() on projection expressions
   - This ensures that projected expressions (e.g., `col + 1`, `upper(name)`) have accurate cardinality estimates

5. **Aggregation constraint propagation** (RelationOp.cpp):
   - Modified Aggregation constructor to accept PlanState parameter
   - Modified `setCostWithGroups()` to accept PlanState
   - Grouping keys get cardinality from the aggregation's result cardinality (number of groups)
   - Aggregate functions get cardinality from exprConstraint() evaluation

6. **VeloxHistory integration** (VeloxHistory.cpp, VeloxHistory.h):
   - Added `setBaseTableValues()` function to update BaseTable column Values from ConstraintMap
   - Modified `findLeafSelectivity()` to always:
     1. Call conjunctsSelectivity() with updateConstraints=true to get constraints
     2. Update BaseTable column Values using setBaseTableValues()
     3. Optionally sample if sampling is enabled
   - This ensures filter-derived constraints (min/max, cardinality, null fractions) are applied to base table columns before planning

7. **Optimization.cpp updates**:
   - Updated all RelationOp construction sites to pass PlanState:
     - Project: added state parameter to constructor calls
     - Aggregation: added state parameter, including planSingleAggregation()
     - Join: added innerFanout parameter and state
     - Filter: already had state
   - Modified PrecomputeProjection::maybeProject() to accept PlanState parameter
   - Updated makeDistinct() to accept PlanState
   - Updated Join::makeCrossJoin() to accept PlanState
   - Ensured PlanStateSaver preserves constraints when exploring alternative join orders

8. **QueryGraphContext lifetime management** (QueryGraphContext.h):
   - Added `registerAny()` template function to take ownership of arbitrary objects (stored as shared_ptr<void>)
   - Added `ownedObjects_` set and `mutex_` for thread-safe lifetime management
   - Used to manage ConstraintMap pointers in Plan objects, ensuring they remain valid throughout optimization

9. **Schema.h/Schema.cpp updates**:
   - Changed Value::cardinality from const to non-const to allow constraint updates
   - Added Value assignment operator that validates type equality before assigning
   - Added Value::toString() for debugging

10. **ToGraph.cpp updates**:
    - Updated all plan construction to pass PlanState to constructors
    - Modified constant deduplication to set min/max for literals to the literal value

11. **FunctionRegistry.h extension**:
    - Added optional `functionConstraint` callback to FunctionMetadata
    - Allows functions to provide custom constraint derivation logic
    - Returns `std::optional<Value>` with refined constraints for the function result

12. **Comprehensive test coverage** (ConstraintsTest.cpp, 347 lines):
    - `scanEquality`: Tests join equality constraints (n_nationkey = r_regionkey), verifies min/max ranges intersect and cardinality is limited to intersection size
    - `aggregateConstraint`: Tests grouping key cardinality propagates to all output columns
    - `projectConstraint`: Tests projected expression cardinality (col+1, col+col) inherits from source columns
    - `outer`: Tests outer join null fraction propagation to optional side (left join, right side becomes nullable with nullFraction ≈ 0.8)
    - `bitwiseAnd`: Tests custom functionConstraint for bitwise_and, verifies min=0 and max=min(arg1.max, arg2.max)

The implementation maintains the invariant that state.constraints contains the most up-to-date Value metadata for all expressions in the current plan state. When exploring alternative join orders or plan structures, PlanStateSaver ensures constraints are properly saved and restored. This enables accurate "what-if" analysis during optimization without polluting the global expression Value metadata.

The constraint propagation integrates seamlessly with the filter selectivity estimation from D88164653:
- Filters call conjunctsSelectivity() which updates state.constraints
- Joins use columnComparisonSelectivity() which updates constraints for join keys
- All downstream operators see refined constraints via value(state, expr)
- Cost estimation at each operator uses the most accurate available cardinality

Future enhancements enabled by this infrastructure:
- Constraint-based partition pruning (skip partitions outside min/max range)
- Dynamic filter pushdown (propagate join-derived constraints to scans)
- Constraint-based empty result detection (min > max indicates zero rows)
- Correlation detection (tracking when columns have matching values)

Differential Revision: D89130357

Author: oerling

axiom/optimizer/Filters.cpp

@@ -127,9 +127,32 @@ float combineSelectivities(
 }
 
 const Value& value(const PlanState& state, ExprCP expr) {
+  auto it = state.constraints.find(expr->id());
+  if (it != state.constraints.end()) {
+    return it->second;
+  }
   return expr->value();
 }
 
+void addConstraint(int32_t exprId, Value value, ConstraintMap& constraints) {
+  // Limit cardinality based on type kind
+  auto typeKind = value.type->kind();
+
+  if (typeKind == velox::TypeKind::BOOLEAN) {
+    // Boolean can have at most 2 distinct values (true/false)
+    value.cardinality = std::min(value.cardinality, 2.0f);
+  } else if (typeKind == velox::TypeKind::TINYINT) {
+    // TINYINT (int8) can have at most 256 distinct values (-128 to 127)
+    value.cardinality = std::min(value.cardinality, 256.0f);
+  } else if (typeKind == velox::TypeKind::SMALLINT) {
+    // SMALLINT (int16) can have at most 65536 distinct values
+    value.cardinality = std::min(value.cardinality, 65536.0f);
+  }
+  // Other types (INTEGER, BIGINT, VARCHAR, etc.) have no practical limit
+
+  constraints.insert_or_assign(exprId, value);
+}
+
 Selectivity comparisonSelectivity(
     const PlanState& state,
     ExprCP expr,
@@ -216,10 +239,15 @@ Selectivity functionSelectivity(
 }
 
 Selectivity conjunctsSelectivity(
-    const PlanState& state,
+    PlanState& state,
     std::span<const ExprCP> conjuncts,
     bool updateConstraints,
     ConstraintMap& newConstraints) {
+  // Update constraints for each conjunct before processing
+  for (auto* conjunct : conjuncts) {
+    exprConstraint(conjunct, state, true);
+  }
+
   std::vector<Selectivity> selectivities;
   selectivities.reserve(conjuncts.size());
 
@@ -271,7 +299,7 @@ Selectivity conjunctsSelectivity(
 }
 
 Selectivity exprSelectivity(
-    const PlanState& state,
+    PlanState& state,
     ExprCP expr,
     bool updateConstraints,
     ConstraintMap& newConstraints) {
@@ -1313,4 +1341,20 @@ Selectivity rangeSelectivity(
   return {0.5 * (1.0 - nullFrac), nullFrac};
 }
 
+/// Declared in namespace to allow calling from debugger.
+std::string constraintsString(ConstraintMap& constraints) {
+  std::stringstream out;
+  for (const auto& pair : constraints) {
+    out << pair.first;
+    if (queryCtx() != nullptr) {
+      auto* expr = queryCtx()->objectAt(pair.first);
+      if (expr != nullptr) {
+        out << " (" << expr->toString() << ")";
+      }
+    }
+    out << " = " << pair.second.toString() << "\n";
+  }
+  return out.str();
+}
+
 } // namespace facebook::axiom::optimizer

axiom/optimizer/Filters.h

@@ -43,6 +43,10 @@ float combineSelectivities(
 
 const Value& value(const PlanState& state, ExprCP expr);
 
+/// Adds a constraint to the constraint map with cardinality limited by type.
+/// For BOOLEAN: max 2, TINYINT: max 256, SMALLINT: max 65536.
+void addConstraint(int32_t exprId, Value value, ConstraintMap& constraints);
+
 Selectivity columnComparisonSelectivity(
     ExprCP left,
     ExprCP right,
@@ -53,13 +57,13 @@ Selectivity columnComparisonSelectivity(
     ConstraintMap& constraints);
 
 Selectivity exprSelectivity(
-    const PlanState& state,
+    PlanState& state,
     ExprCP expr,
     bool updateConstraints,
     ConstraintMap& newConstraints);
 
 Selectivity conjunctsSelectivity(
-    const PlanState& state,
+    PlanState& state,
     std::span<const ExprCP> conjuncts,
     bool updateConstraints,
     ConstraintMap& newConstraints);

axiom/optimizer/FunctionRegistry.h

@@ -20,6 +20,9 @@
 
 namespace facebook::axiom::optimizer {
 
+struct PlanState;
+struct Value;
+
 /// A bit set that qualifies an Expr. Represents which functions/kinds
 /// of functions are found inside the children of an Expr.
 class FunctionSet {
@@ -182,6 +185,10 @@ struct FunctionMetadata {
       const logical_plan::CallExpr* call,
       std::vector<PathCP>& paths)>
       explode;
+
+  /// Function to compute derived constraints for function calls.
+  std::function<std::optional<Value>(ExprCP, PlanState& state)>
+      functionConstraint;
 };
 
 using FunctionMetadataCP = const FunctionMetadata*;

axiom/optimizer/JoinSample.cpp

@@ -142,15 +142,20 @@ std::shared_ptr<runner::Runner> prepareSampleRunner(
       make<Call>(toName(kHashMix), bigintValue(), hashes, FunctionSet{});
 
   ColumnCP hashColumn = make<Column>(toName("hash"), nullptr, hash->value());
+
+  // Create temporary PlanState for Project and Filter constructors
+  PlanState tempState(*queryCtx()->optimization(), nullptr);
+
   RelationOpPtr project = make<Project>(
-      scan, ExprVector{hash}, ColumnVector{hashColumn}, /*redundant=*/false);
+      scan,
+      ExprVector{hash},
+      ColumnVector{hashColumn},
+      /*redundant=*/false,
+      tempState);
 
   // (hash % mod) < lim
   ExprCP filterExpr = makeCall(
       kSample, velox::BOOLEAN(), hashColumn, bigintLit(mod), bigintLit(lim));
-
-  // Create temporary PlanState for Filter constructor
-  PlanState tempState(*queryCtx()->optimization(), nullptr);
   RelationOpPtr filter =
       make<Filter>(tempState, project, ExprVector{filterExpr});