Track Value constraints through plan construction (facebookincubator#708)

Summary: Pull Request resolved: facebookincubator#708

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});