Add min_field_depth_for_optimization option to SelectOptimizationOptions.

PiperOrigin-RevId: 1000374223

Author: jnthntatum

extensions/select_optimization.cc

@@ -400,8 +400,9 @@ absl::StatusOr<std::vector<SelectQualifier>> SelectInstructionsFromCall(
 
 class RewriterImpl : public AstRewriterBase {
  public:
-  RewriterImpl(const Ast& ast, PlannerContext& planner_context)
-      : ast_(ast), planner_context_(planner_context) {}
+  RewriterImpl(const Ast& ast, PlannerContext& planner_context,
+               const SelectOptimizationOptions& options)
+      : ast_(ast), planner_context_(planner_context), options_(options) {}
 
   void PreVisitExpr(const Expr& expr) override { path_.push_back(&expr); }
 
@@ -474,6 +475,10 @@ class RewriterImpl : public AstRewriterBase {
     }
 
     SelectPath path = GetSelectPath(&expr);
+    if (path.select_instructions.size() <
+        options_.min_field_depth_for_optimization) {
+      return false;
+    }
 
     // generate the new cel.attribute call.
     absl::string_view fn = path.test_only ? kCelHasField : kCelAttribute;
@@ -556,6 +561,7 @@ class RewriterImpl : public AstRewriterBase {
 
   const Ast& ast_;
   PlannerContext& planner_context_;
+  SelectOptimizationOptions options_;
   // ids of potentially optimizeable expr nodes.
   absl::flat_hash_map<const Expr*, QualifierInstruction> candidates_;
   std::vector<const Expr*> path_;
@@ -858,7 +864,7 @@ google::api::expr::runtime::FlatExprBuilder* GetFlatExprBuilder(
 
 absl::Status SelectOptimizationAstUpdater::UpdateAst(PlannerContext& context,
                                                      Ast& ast) const {
-  RewriterImpl rewriter(ast, context);
+  RewriterImpl rewriter(ast, context, options_);
   AstRewrite(ast.mutable_root_expr(), rewriter);
   return rewriter.GetProgressStatus();
 }
@@ -880,7 +886,7 @@ absl::Status EnableSelectOptimization(
   }
 
   flat_expr_builder->AddAstTransform(
-      std::make_unique<SelectOptimizationAstUpdater>());
+      std::make_unique<SelectOptimizationAstUpdater>(options));
   // Add overloads for select optimization signature.
   // These are never bound, only used to prevent the builder from failing on
   // the overloads check.

extensions/select_optimization.h

@@ -35,6 +35,15 @@ struct SelectOptimizationOptions {
   // unimplemented for a given StructType. This option is exposed for testing or
   // to more closely match behavior of unoptimized expressions.
   bool force_fallback_implementation = false;
+
+  // The minimum depth of a select chain to be considered for optimization.
+  // For example, if this is set to 2, then `a.b.c` will be optimized, but
+  // `a.b` will not be.
+  //
+  // This is applied in the pre-pass phase that rewrites the AST. If the
+  // optimization is already applied by a compiler / static optimizer, this
+  // option has no effect.
+  int min_field_depth_for_optimization = 0;
 };
 
 // Enable select optimization on the given RuntimeBuilder, replacing long
@@ -77,10 +86,15 @@ absl::Status EnableSelectOptimization(
 class SelectOptimizationAstUpdater
     : public google::api::expr::runtime::AstTransform {
  public:
-  SelectOptimizationAstUpdater() = default;
+  explicit SelectOptimizationAstUpdater(
+      const SelectOptimizationOptions& options = {})
+      : options_(options) {}
 
   absl::Status UpdateAst(google::api::expr::runtime::PlannerContext& context,
                          cel::Ast& ast) const override;
+
+ private:
+  SelectOptimizationOptions options_;
 };
 
 google::api::expr::runtime::ProgramOptimizerFactory

extensions/select_optimization_test.cc

@@ -528,6 +528,47 @@ TEST_F(SelectOptimizationTest, AstTransformComplexSelect) {
       ElementsAre(SelectFieldEntry(1, "child"), SelectFieldEntry(1, "child")));
 }
 
+TEST_F(SelectOptimizationTest, AstTransformMinFieldDepth) {
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<cel::Ast> ast,
+      CompileForTestCase(
+          "((false)? a.child.child : b.child).child.payload.single_int64"));
+
+  SelectOptimizationOptions options;
+  options.min_field_depth_for_optimization = 2;
+  SelectOptimizationAstUpdater updater(options);
+  EXPECT_THAT(updater.UpdateAst(context_, *ast), IsOk());
+
+  const CallExpr& attr_call = ast->root_expr().call_expr();
+  EXPECT_EQ(attr_call.function(), "cel.@attribute");
+
+  ASSERT_THAT(attr_call.args(), SizeIs(2));
+
+  EXPECT_THAT(
+      attr_call.args()[1].list_expr().elements(),
+      ElementsAre(SelectFieldEntry(1, "child"), SelectFieldEntry(2, "payload"),
+                  SelectFieldEntry(2, "single_int64")));
+
+  const Expr& operand = attr_call.args()[0];
+
+  EXPECT_EQ(operand.call_expr().function(), cel::builtin::kTernary);
+  ASSERT_THAT(operand.call_expr().args(), SizeIs(3));
+
+  const Expr& true_branch = operand.call_expr().args()[1];
+
+  EXPECT_EQ(true_branch.call_expr().function(), "cel.@attribute");
+  ASSERT_THAT(true_branch.call_expr().args(), SizeIs(2));
+
+  EXPECT_THAT(
+      true_branch.call_expr().args()[1].list_expr().elements(),
+      ElementsAre(SelectFieldEntry(1, "child"), SelectFieldEntry(1, "child")));
+
+  const Expr& false_branch = operand.call_expr().args()[2];
+  EXPECT_TRUE(false_branch.has_select_expr());
+  EXPECT_EQ(false_branch.select_expr().field(), "child");
+  EXPECT_EQ(false_branch.select_expr().operand().ident_expr().name(), "b");
+}
+
 TEST_F(SelectOptimizationTest, AstTransformMapIndexTraversal) {
   // nested_test_all_types.payload.map_string_message['$not_a_field'].bb
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<cel::Ast> ast,