Track Value constraints through plan construction (facebookincubator#708)

Summary: Pull Request resolved: facebookincubator#708

Differential Revision: D89130357

Author: oerling

axiom/optimizer/Filters.cpp

@@ -127,7 +127,11 @@ float combineSelectivities(
 }
 
 const Value& value(const PlanState& state, ExprCP expr) {
-  return expr->value();
+  auto it = state.constraints.find(expr->id());
+  if (it != state.constraints.end()) {
+    return it->second;
+  }
+    return expr->value();
 }
 
 Selectivity comparisonSelectivity(

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/Optimization.cpp

@@ -19,6 +19,7 @@
 #include <iostream>
 #include <utility>
 #include "axiom/optimizer/DerivedTablePrinter.h"
+#include "axiom/optimizer/Filters.h"
 #include "axiom/optimizer/Plan.h"
 #include "axiom/optimizer/PrecomputeProjection.h"
 #include "axiom/optimizer/VeloxHistory.h"
@@ -820,7 +821,7 @@ void Optimization::addPostprocess(
     VELOX_DCHECK(!dt->hasLimit());
     PrecomputeProjection precompute{plan, dt, /*projectAllInputs=*/false};
     auto writeColumns = precompute.toColumns(dt->write->columnExprs());
-    plan = std::move(precompute).maybeProject();
+    plan = std::move(precompute).maybeProject(state);
     state.addCost(*plan);
 
     plan = repartitionForWrite(plan, state);
@@ -868,7 +869,8 @@ void Optimization::addPostprocess(
         maybeDropProject(plan),
         usedExprs,
         usedColumns,
-        Project::isRedundant(plan, usedExprs, usedColumns));
+        Project::isRedundant(plan, usedExprs, usedColumns),
+        state);
   }
 
   if (!dt->hasOrderBy() && dt->limit > kMaxLimitBeforeProject) {
@@ -915,7 +917,8 @@ AggregateVector flattenAggregates(
 // static
 RelationOpPtr Optimization::planSingleAggregation(
     DerivedTableCP dt,
-    RelationOpPtr& input) {
+    RelationOpPtr& input,
+    PlanState& state) {
   const auto* aggPlan = dt->aggregation;
 
   PrecomputeProjection precompute(input, dt, /*projectAllInputs=*/false);
@@ -924,11 +927,12 @@ RelationOpPtr Optimization::planSingleAggregation(
   auto aggregates = flattenAggregates(aggPlan->aggregates(), precompute);
 
   return make<Aggregation>(
-      std::move(precompute).maybeProject(),
+      std::move(precompute).maybeProject(state),
       std::move(groupingKeys),
       std::move(aggregates),
       velox::core::AggregationNode::Step::kSingle,
-      aggPlan->columns());
+      aggPlan->columns(),
+      state);
 }
 
 void Optimization::addAggregation(
@@ -942,7 +946,7 @@ void Optimization::addAggregation(
       precompute.toColumns(aggPlan->groupingKeys(), &aggPlan->columns());
   auto aggregates = flattenAggregates(aggPlan->aggregates(), precompute);
 
-  plan = std::move(precompute).maybeProject();
+  plan = std::move(precompute).maybeProject(state);
   state.placed.add(aggPlan);
 
   if (isSingleWorker_ && isSingleDriver_) {
@@ -951,7 +955,8 @@ void Optimization::addAggregation(
         std::move(groupingKeys),
         std::move(aggregates),
         velox::core::AggregationNode::Step::kSingle,
-        aggPlan->columns());
+        aggPlan->columns(),
+        state);
 
     state.addCost(*singleAgg);
     plan = singleAgg;
@@ -972,7 +977,8 @@ void Optimization::addAggregation(
       groupingKeys,
       aggregates,
       velox::core::AggregationNode::Step::kPartial,
-      aggPlan->intermediateColumns());
+      aggPlan->intermediateColumns(),
+      state);
 
   PlanCost splitAggCost;
   splitAggCost.add(*partialAgg);
@@ -991,7 +997,8 @@ void Optimization::addAggregation(
       std::move(finalGroupingKeys),
       aggregates,
       velox::core::AggregationNode::Step::kFinal,
-      aggPlan->columns());
+      aggPlan->columns(),
+      state);
   splitAggCost.add(*splitAggPlan);
 
   if (numKeys == 0 || options_.alwaysPlanPartialAggregation) {
@@ -1009,7 +1016,8 @@ void Optimization::addAggregation(
       groupingKeys,
       aggregates,
       velox::core::AggregationNode::Step::kSingle,
-      aggPlan->columns());
+      aggPlan->columns(),
+      state);
   singleAggCost.add(*singleAgg);
 
   if (singleAggCost.cost < splitAggCost.cost) {
@@ -1043,7 +1051,7 @@ void Optimization::addOrderBy(
   }
 
   auto* orderBy = make<OrderBy>(
-      std::move(precompute).maybeProject(),
+      std::move(precompute).maybeProject(state),
       std::move(orderKeys),
       dt->orderTypes,
       dt->limit,
@@ -1136,9 +1144,9 @@ struct ProjectionBuilder {
     exprs.emplace_back(expr);
   }
 
-  RelationOp* build(RelationOp* input) {
+  RelationOp* build(RelationOp* input, PlanState& state) {
     return make<Project>(
-        input, exprs, columns, Project::isRedundant(input, exprs, columns));
+        input, exprs, columns, Project::isRedundant(input, exprs, columns), state);
   }
 
   ColumnVector inputColumns() const {
@@ -1229,6 +1237,42 @@ void tryOptimizeSemiProject(
 }
 } // namespace
 
+void Optimization::addJoinConstraint(
+    ExprCP left,
+    ExprCP right,
+    bool leftOptional,
+    bool rightOptional,
+    PlanState& state) {
+  // Get the values for left and right expressions
+  Value leftValue = value(state, left);
+  Value rightValue = value(state, right);
+
+  // Set nullability based on optionality flags
+  if (leftOptional) {
+    leftValue.nullFraction = std::max(leftValue.nullFraction, 0.01f);
+  }
+  if (rightOptional) {
+    rightValue.nullFraction = std::max(rightValue.nullFraction, 0.01f);
+  }
+
+  // Call columnComparisonSelectivity with updateConstraints=true
+  columnComparisonSelectivity(
+      left, right, leftValue, rightValue, toName("eq"), true, state.constraints);
+}
+
+void Optimization::addJoinConstraints(
+    const ExprVector& left,
+    const ExprVector& right,
+    bool leftOptional,
+    bool rightOptional,
+    PlanState& state) {
+  VELOX_CHECK_EQ(left.size(), right.size(), "Join key vectors must have same size");
+
+  for (size_t i = 0; i < left.size(); ++i) {
+    addJoinConstraint(left[i], right[i], leftOptional, rightOptional, state);
+  }
+}
+
 void Optimization::joinByHash(
     const RelationOpPtr& plan,
     const JoinCandidate& candidate,
@@ -1330,7 +1374,7 @@ void Optimization::joinByHash(
 
   PrecomputeProjection precomputeBuild(buildInput, state.dt);
   auto buildKeys = precomputeBuild.toColumns(build.keys);
-  buildInput = std::move(precomputeBuild).maybeProject();
+  buildInput = std::move(precomputeBuild).maybeProject(buildState);
 
   auto* buildOp = make<HashBuild>(buildInput, build.keys, buildPlan);
   buildState.addCost(*buildOp);
@@ -1396,7 +1440,13 @@ void Optimization::joinByHash(
 
   PrecomputeProjection precomputeProbe(probeInput, state.dt);
   auto probeKeys = precomputeProbe.toColumns(probe.keys);
-  probeInput = std::move(precomputeProbe).maybeProject();
+  probeInput = std::move(precomputeProbe).maybeProject(state);
+
+  // Add join constraints for equi-join keys (except for anti joins)
+  if (joinType != velox::core::JoinType::kAnti) {
+    bool rightOptional = (joinType == velox::core::JoinType::kLeft);
+    addJoinConstraints(probe.keys, build.keys, false, rightOptional, state);
+  }
 
   RelationOp* join = make<Join>(
       JoinMethod::kHash,
@@ -1413,7 +1463,7 @@ void Optimization::joinByHash(
   state.cost.cost += buildState.cost.cost;
 
   if (needsProjection) {
-    join = projectionBuilder.build(join);
+    join = projectionBuilder.build(join, state);
   }
 
   state.addNextJoin(&candidate, join, toTry);
@@ -1475,7 +1525,7 @@ void Optimization::joinByHashRight(
 
   PrecomputeProjection precomputeBuild(buildInput, state.dt);
   auto buildKeys = precomputeBuild.toColumns(build.keys);
-  buildInput = std::move(precomputeBuild).maybeProject();
+  buildInput = std::move(precomputeBuild).maybeProject(state);
 
   auto* buildOp = make<HashBuild>(buildInput, build.keys, nullptr);
   state.addCost(*buildOp);
@@ -1554,7 +1604,12 @@ void Optimization::joinByHashRight(
 
   PrecomputeProjection precomputeProbe(probeInput, state.dt);
   auto probeKeys = precomputeProbe.toColumns(probe.keys);
-  probeInput = std::move(precomputeProbe).maybeProject();
+  probeInput = std::move(precomputeProbe).maybeProject(state);
+
+  // Add join constraints for equi-join keys
+  // In joinByHashRight, if join type is right outer, leftOptional is true
+  bool leftOptional = (rightJoinType == velox::core::JoinType::kRight);
+  addJoinConstraints(probe.keys, build.keys, leftOptional, false, state);
 
   RelationOp* join = make<Join>(
       JoinMethod::kHash,
@@ -1569,7 +1624,7 @@ void Optimization::joinByHashRight(
   state.addCost(*join);
 
   if (needsProjection) {
-    join = projectionBuilder.build(join);
+    join = projectionBuilder.build(join, state);
   }
 
   state.addNextJoin(&candidate, join, toTry);
@@ -1613,7 +1668,7 @@ void Optimization::crossJoinUnnest(
     // because we can have multiple unnest joins in single JoinCandidate.
 
     auto unnestColumns = precompute.toColumns(unnestExprs);
-    plan = std::move(precompute).maybeProject();
+    plan = std::move(precompute).maybeProject(state);
 
     plan = make<Unnest>(
         std::move(plan),
@@ -1853,7 +1908,9 @@ ColumnVector indexColumns(
   return result;
 }
 
-RelationOpPtr makeDistinct(const RelationOpPtr& input) {
+RelationOpPtr makeDistinct(
+    const RelationOpPtr& input,
+    const PlanState& state) {
   ExprVector groupingKeys;
   for (const auto& column : input->columns()) {
     groupingKeys.push_back(column);
@@ -1864,7 +1921,8 @@ RelationOpPtr makeDistinct(const RelationOpPtr& input) {
       groupingKeys,
       AggregateVector{},
       velox::core::AggregationNode::Step::kSingle,
-      input->columns());
+      input->columns(),
+      state);
 }
 
 Distribution somePartition(const RelationOpPtrVector& inputs) {
@@ -2104,7 +2162,7 @@ PlanP Optimization::makeUnionPlan(
     RelationOpPtr result = make<UnionAll>(inputs);
     Aggregation* distinct = nullptr;
     if (isDistinct) {
-      result = makeDistinct(result);
+      result = makeDistinct(result, inputStates[0]);
       distinct = result->as<Aggregation>();
     }
     return unionPlan(inputStates, result, distinct);
@@ -2135,7 +2193,7 @@ PlanP Optimization::makeUnionPlan(
   RelationOpPtr result = make<UnionAll>(inputs);
   Aggregation* distinct = nullptr;
   if (isDistinct) {
-    result = makeDistinct(result);
+    result = makeDistinct(result, inputStates[0]);
     distinct = result->as<Aggregation>();
   }
   return unionPlan(inputStates, result, distinct);

axiom/optimizer/Optimization.h

@@ -121,7 +121,8 @@ class Optimization {
   /// @param dt Derived table with an aggregation.
   static RelationOpPtr planSingleAggregation(
       DerivedTableCP dt,
-      RelationOpPtr& input);
+      RelationOpPtr& input,
+      PlanState& state);
 
   const std::shared_ptr<velox::core::QueryCtx>& veloxQueryCtx() const {
     return veloxQueryCtx_;
@@ -302,6 +303,22 @@ class Optimization {
       PlanState& state,
       std::vector<NextJoin>& toTry);
 
+  // Adds join constraints for a single pair of join keys.
+  void addJoinConstraint(
+      ExprCP left,
+      ExprCP right,
+      bool leftOptional,
+      bool rightOptional,
+      PlanState& state);
+
+  // Adds join constraints for vectors of join keys.
+  void addJoinConstraints(
+      const ExprVector& left,
+      const ExprVector& right,
+      bool leftOptional,
+      bool rightOptional,
+      PlanState& state);
+
   void crossJoin(
       const RelationOpPtr& plan,
       const JoinCandidate& candidate,