Whole function subfield tracking

Differential Revision: D89635456

Author: Orri Erling

axiom/optimizer/FunctionRegistry.h

@@ -121,7 +121,8 @@ class Call;
 struct FunctionMetadata {
   bool processSubfields() const {
     return subfieldArg.has_value() || !fieldIndexForArg.empty() ||
-        isArrayConstructor || isMapConstructor || valuePathToArgPath;
+        isArrayConstructor || isMapConstructor || valuePathToArgPath ||
+        explode || expandFunction || !lambdas.empty();
   }
 
   const LambdaInfo* lambdaInfo(int32_t index) const {
@@ -186,6 +187,16 @@ struct FunctionMetadata {
       std::vector<PathCP>& paths)>
       explode;
 
+  /// Hook for rewriting a call to a function. In the case of a
+  /// complex type function with subfield related metadata,
+  /// 'logicalExplode' is used if there are only getters over the
+  /// function. For functions with subfield related metadata options,
+  /// 'expandFunction is used only if the function is accessed as a
+  /// whole. If returns non-nullptr, the returned expression is used
+  /// in the place of the function.
+  std::function<logical_plan::ExprPtr(const logical_plan::CallExpr*)>
+      expandFunction;
+
   /// Function to compute derived constraints for function calls.
   std::function<std::optional<Value>(ExprCP, PlanState& state)>
       functionConstraint;

axiom/optimizer/SubfieldTracker.cpp

@@ -24,6 +24,25 @@ namespace lp = facebook::axiom::logical_plan;
 
 namespace facebook::axiom::optimizer {
 
+lp::ExprPtr ExpandFunctionCache::getOrExpand(
+    const lp::Expr* expr,
+    const std::function<lp::ExprPtr(const lp::CallExpr*)>& expandFunction) {
+  auto it = cache_.find(expr);
+  if (it != cache_.end()) {
+    return it->second;
+  }
+
+  auto* call = dynamic_cast<const lp::CallExpr*>(expr);
+  VELOX_CHECK_NOT_NULL(call, "Expression must be a CallExpr");
+
+  auto result = expandFunction(call);
+  // Only cache non-null results to avoid caching failures
+  if (result) {
+    cache_[expr] = result;
+  }
+  return result;
+}
+
 SubfieldTracker::SubfieldTracker(
     std::function<logical_plan::ConstantExprPtr(const logical_plan::ExprPtr&)>
         tryFoldConstant)
@@ -428,7 +447,17 @@ void SubfieldTracker::markSubfields(
     // If the function is some kind of constructor, like
     // make_row_from_map or make_named_row, then a path over it
     // selects one argument. If there is no path, all arguments are
-    // implicitly accessed.
+    // implicitly accessed. If the whole value of make_row_from_map is accessed,
+    // this does not yet access the whole map but just the keys listed in
+    // make_row_from_map.
+    if (metadata->expandFunction && steps.empty()) {
+      auto newExpr =
+          expandFunctionCache_.getOrExpand(call, metadata->expandFunction);
+      if (newExpr) {
+        markSubfields(newExpr, steps, isControl, context);
+        return;
+      }
+    }
     if (metadata->valuePathToArgPath && !steps.empty()) {
       auto pair = metadata->valuePathToArgPath(steps, *call);
       markSubfields(expr->inputAt(pair.second), pair.first, isControl, context);

axiom/optimizer/SubfieldTracker.h

@@ -21,6 +21,42 @@
 
 namespace facebook::axiom::optimizer {
 
+/// Cache for expandFunction results. Maps from lp::Expr* to expanded ExprPtr.
+/// Uses pointer-based hashing and comparison.
+struct ExpandFunctionCache {
+  /// Checks if the expression is in the cache and returns the cached result.
+  /// If not in cache, calls expandFunction, caches the result, and returns it.
+  /// @param expr The expression to expand (must be a CallExpr)
+  /// @param expandFunction The expansion function to call if not cached
+  /// @return The expanded expression, or nullptr if expansion returns nullptr
+  logical_plan::ExprPtr getOrExpand(
+      const logical_plan::Expr* expr,
+      const std::function<logical_plan::ExprPtr(const logical_plan::CallExpr*)>&
+          expandFunction);
+
+ private:
+  struct PointerHash {
+    size_t operator()(const logical_plan::Expr* ptr) const {
+      return std::hash<const logical_plan::Expr*>()(ptr);
+    }
+  };
+
+  struct PointerEqual {
+    bool operator()(
+        const logical_plan::Expr* lhs,
+        const logical_plan::Expr* rhs) const {
+      return lhs == rhs;
+    }
+  };
+
+  folly::F14FastMap<
+      const logical_plan::Expr*,
+      logical_plan::ExprPtr,
+      PointerHash,
+      PointerEqual>
+      cache_;
+};
+
 /// Set of accessed subfields given ordinal of output column or function
 /// argument.
 struct ResultAccess {
@@ -74,10 +110,14 @@ class SubfieldTracker {
 
   /// Goes over the local plan and collects all accessed columns and subfields.
   /// Reports 'control' and 'payload' columns and subfields separately.
-  std::pair<PlanSubfields, PlanSubfields> markAll(
+  /// Returns the subfields and the expandFunction cache that can be reused.
+  std::tuple<PlanSubfields, PlanSubfields, ExpandFunctionCache> markAll(
       const logical_plan::LogicalPlanNode& node) && {
     markAllSubfields(node, {});
-    return {controlSubfields_, payloadSubfields_};
+    return {
+        std::move(controlSubfields_),
+        std::move(payloadSubfields_),
+        std::move(expandFunctionCache_)};
   }
 
   // if 'step' applied to result of the function of 'metadata'
@@ -147,6 +187,8 @@ class SubfieldTracker {
 
   PlanSubfields controlSubfields_;
   PlanSubfields payloadSubfields_;
+
+  ExpandFunctionCache expandFunctionCache_;
 };
 
 } // namespace facebook::axiom::optimizer

axiom/optimizer/ToGraph.cpp

@@ -1024,9 +1024,10 @@ ExprCP ToGraph::translateExpr(const lp::ExprPtr& expr) {
 
   const auto* call = expr->isCall() ? expr->as<lp::CallExpr>() : nullptr;
   std::string callName;
+  const FunctionMetadata* metadata = nullptr;
   if (call) {
     callName = velox::exec::sanitizeName(call->name());
-    auto* metadata = functionMetadata(callName);
+    metadata = functionMetadata(callName);
     if (metadata && metadata->processSubfields()) {
       auto translated = translateSubfieldFunction(call, metadata);
       if (translated.has_value()) {
@@ -1035,6 +1036,14 @@ ExprCP ToGraph::translateExpr(const lp::ExprPtr& expr) {
     }
   }
 
+  if (metadata && metadata->expandFunction) {
+    auto newExpr =
+        expandFunctionCache_.getOrExpand(call, metadata->expandFunction);
+    if (newExpr) {
+      return translateExpr(newExpr);
+    }
+  }
+
   const auto* specialForm =
       expr->isSpecialForm() ? expr->as<lp::SpecialFormExpr>() : nullptr;
 
@@ -2492,7 +2501,7 @@ void ToGraph::translateUnion(const lp::SetNode& set) {
 }
 
 DerivedTableP ToGraph::makeQueryGraph(const lp::LogicalPlanNode& logicalPlan) {
-  std::tie(controlSubfields_, payloadSubfields_) =
+  std::tie(controlSubfields_, payloadSubfields_, expandFunctionCache_) =
       SubfieldTracker([&](const auto& expr) {
         return tryFoldConstant(expr);
       }).markAll(logicalPlan);

axiom/optimizer/ToGraph.h

@@ -419,6 +419,9 @@ class ToGraph {
   // Column and subfield info for items that only affect column values.
   PlanSubfields payloadSubfields_;
 
+  // Cache for expandFunction results, populated during SubfieldTracker::markAll
+  ExpandFunctionCache expandFunctionCache_;
+
   /// Expressions corresponding to skyline paths over a subfield decomposable
   /// function.
   folly::F14FastMap<const logical_plan::CallExpr*, SubfieldProjections>

axiom/optimizer/tests/Genies.cpp

@@ -61,9 +61,46 @@ VELOX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
     exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
     std::make_unique<GenieFunction>());
 
+class IdentityFunction : public exec::VectorFunction {
+ public:
+  void apply(
+      const SelectivityVector& rows,
+      std::vector<VectorPtr>& args,
+      const TypePtr& /* outputType */,
+      exec::EvalCtx& context,
+      VectorPtr& result) const override {
+    VELOX_CHECK_EQ(args.size(), 1);
+
+    auto& input = args[0];
+    if (!result) {
+      // If no result vector provided, just return the input vector directly
+      result = input;
+    } else {
+      // If result vector is provided, copy input to result
+      BaseVector::ensureWritable(rows, input->type(), context.pool(), result);
+      result->copy(input.get(), rows, nullptr);
+    }
+  }
+
+  static std::vector<std::shared_ptr<exec::FunctionSignature>> signatures() {
+    // T -> T (accepts any type and returns the same type)
+    return {exec::FunctionSignatureBuilder()
+                .typeVariable("T")
+                .returnType("T")
+                .argumentType("T")
+                .build()};
+  }
+};
+
+VELOX_DECLARE_VECTOR_FUNCTION(
+    udf_identity,
+    IdentityFunction::signatures(),
+    std::make_unique<IdentityFunction>());
+
 void registerGenieUdfs() {
   VELOX_REGISTER_VECTOR_FUNCTION(udf_genie, "genie");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_genie, "exploding_genie");
+  VELOX_REGISTER_VECTOR_FUNCTION(udf_identity, "identity");
 }
 
 } // namespace facebook::axiom::optimizer::test

axiom/optimizer/tests/SubfieldTest.cpp

@@ -674,6 +674,10 @@ TEST_P(SubfieldTest, overAggregation) {
 }
 
 TEST_P(SubfieldTest, blackbox) {
+  registerGenieUdfs();
+  if (GetParam() == 3) {
+    optimizerOptions_.allMapsAsStruct = true;
+  }
   auto data = makeRowVector(
       {"id", "m"},
       {
@@ -686,7 +690,8 @@ TEST_P(SubfieldTest, blackbox) {
 
   lp::PlanBuilder::Context ctx(kHiveConnectorId);
   ctx.hook = [](const auto& name, const auto& args) -> lp::ExprPtr {
-    if (name == "map_row_from_map") {
+    if (name == "map_row_from_map" || name == "make_row_from_map" ||
+        name == "padded_make_row_from_map") {
       VELOX_CHECK(args.at(2)->isConstant());
       auto names = args.at(2)
                        ->template as<lp::ConstantExpr>()
@@ -723,6 +728,82 @@ TEST_P(SubfieldTest, blackbox) {
           .build();
 
   ASSERT_NO_THROW(toSingleNodePlan(logicalPlan));
+
+  logicalPlan =
+      lp::PlanBuilder(ctx)
+          .tableScan("t")
+          .project(
+              {"make_row_from_map(m, array[1, 2, 3], array['f1', 'f2', 'f3']) as m"})
+          .build();
+
+  auto plan = toSingleNodePlan(logicalPlan);
+
+  verifyRequiredSubfields(
+      plan, {{"m", {subfield("1"), subfield("2"), subfield("3")}}});
+
+  if (GetParam() == 1) {
+    auto matcher =
+        core::PlanMatcherBuilder()
+            .tableScan()
+            .project(
+                {"row_constructor(subscript(m_4,1),subscript(m_4,2),subscript(m_4,3))"})
+            .build();
+
+    ASSERT_TRUE(matcher->match(plan));
+  } else {
+    auto matcher =
+        core::PlanMatcherBuilder().tableScan().project().project().build();
+
+    ASSERT_TRUE(matcher->match(plan));
+  }
+
+  logicalPlan =
+      lp::PlanBuilder(ctx, true)
+          .tableScan("t")
+          .project(
+              {"identity(make_row_from_map(m, array[1, 2, 3], array['f1', 'f2', 'f3'])) as m"})
+          .project(
+              {"if (m.f1 < 0, ceil(m.f1), floor(m.f1)) as f1b",
+               "if (m.f2 < 0, 1 + floor(m.f2), ceil(m.f2) + 1) as f2b"})
+          .project({"f1b + 1 as f1b1", "f1b * 2 as f1b2b", "f2b * 3 as f2b3"})
+          .build();
+
+  // Enable parallel project for the remainder of this test. This is reset in
+  // SetUp().
+  optimizerOptions_.parallelProjectWidth = 2;
+  plan = toSingleNodePlan(logicalPlan);
+
+  verifyRequiredSubfields(
+      plan, {{"m", {subfield("1"), subfield("2"), subfield("3")}}});
+
+  if (GetParam() == 1) {
+    auto matcher =
+        core::PlanMatcherBuilder()
+            .tableScan()
+            .parallelProject()
+            .parallelProject(
+                {"identity(row_constructor(subscript(m_7,1),subscript(m_7,2),subscript(m_7,3)))"})
+            .parallelProject()
+            .parallelProject()
+            .project()
+            .build();
+
+    ASSERT_TRUE(matcher->match(plan));
+  }
+
+  logicalPlan =
+      lp::PlanBuilder(ctx, true)
+          .tableScan("t")
+          .project(
+              {"make_row_from_map(m, array[1, 2, 3], array['f1', 'f2', 'f3']) as m"})
+          .project(
+              {"if (coalesce(m.f1, 1::REAL) < 0, ceil(coalesce(m.f1, 1::REAL)), floor(coalesce(m.f1, 1::REAL))) as f1b",
+               "if (coalesce(m.f2, 1::REAL) < 0, 1 + floor(coalesce(m.f2, 2::real)), ceil(coalesce(m.f2, 3::REAL)) + 1) as f2b"})
+          .project({"f1b + 1 as f1b1", "f1b * 2 as f1b2b", "f2b * 3 as f2b3"})
+          .build();
+
+  plan = toSingleNodePlan(logicalPlan);
+  std::cout << plan->toString(true, true);
 }
 
 VELOX_INSTANTIATE_TEST_SUITE_P(