diff --git a/extensions/select_optimization.cc b/extensions/select_optimization.cc index 5f0512011..019812438 100644 --- a/extensions/select_optimization.cc +++ b/extensions/select_optimization.cc @@ -400,8 +400,9 @@ absl::StatusOr> 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 candidates_; std::vector 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()); + std::make_unique(options)); // Add overloads for select optimization signature. // These are never bound, only used to prevent the builder from failing on // the overloads check. diff --git a/extensions/select_optimization.h b/extensions/select_optimization.h index de412afe9..8beeb1bed 100644 --- a/extensions/select_optimization.h +++ b/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 diff --git a/extensions/select_optimization_test.cc b/extensions/select_optimization_test.cc index ea5512976..227b902e8 100644 --- a/extensions/select_optimization_test.cc +++ b/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 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 ast,