Add compiler optimization to remove px.contains calls with empty strings (#1679)

Author: ddelnano

src/carnot/planner/compiler/optimizer/BUILD.bazel

@@ -82,3 +82,13 @@ pl_cc_test(
         "//src/carnot/udf_exporter:cc_library",
     ],
 )
+
+pl_cc_test(
+    name = "prune_unused_contains_rule_test",
+    srcs = ["prune_unused_contains_rule_test.cc"],
+    deps = [
+        ":cc_library",
+        "//src/carnot/planner/compiler:test_utils",
+        "//src/carnot/udf_exporter:cc_library",
+    ],
+)

src/carnot/planner/compiler/optimizer/optimizer.h

@@ -26,6 +26,7 @@
 #include "src/carnot/planner/compiler/optimizer/merge_nodes_rule.h"
 #include "src/carnot/planner/compiler/optimizer/prune_unconnected_operators_rule.h"
 #include "src/carnot/planner/compiler/optimizer/prune_unused_columns_rule.h"
+#include "src/carnot/planner/compiler/optimizer/prune_unused_contains_rule.h"
 #include "src/carnot/planner/compiler_state/compiler_state.h"
 #include "src/carnot/planner/compiler_state/registry_info.h"
 #include "src/carnot/planner/ir/ir.h"
@@ -62,10 +63,16 @@ class Optimizer : public RuleExecutor<IR> {
     prune_unused_columns->AddRule<PruneUnusedColumnsRule>();
   }
 
+  void CreatePruneUnusedContainsBatch() {
+    RuleBatch* prune_unused_columns = CreateRuleBatch<DoOnce>("PruneUnusedContains");
+    prune_unused_columns->AddRule<PruneUnusedContainsRule>();
+  }
+
   Status Init() {
     CreatePruneUnconnectedOpsBatch();
     CreateMergeNodesBatch();
     CreatePruneUnusedColumnsBatch();
+    CreatePruneUnusedContainsBatch();
     return Status::OK();
   }
 

src/carnot/planner/compiler/optimizer/prune_unused_contains_rule.cc

@@ -0,0 +1,88 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "src/carnot/planner/compiler/optimizer/prune_unused_contains_rule.h"
+
+#include <algorithm>
+#include <queue>
+
+namespace px {
+namespace carnot {
+namespace planner {
+namespace compiler {
+
+StatusOr<bool> PruneUnusedContainsRule::RemoveMatchingFilter(IRNode* ir_node) {
+  auto ir_graph = ir_node->graph();
+  auto node_id = ir_node->id();
+  if (!Match(ir_node, Filter())) return false;
+
+  FilterIR* filter = static_cast<FilterIR*>(ir_node);
+  ExpressionIR* expr = filter->filter_expr();
+
+  if (!Match(expr, Func("contains"))) return false;
+
+  FuncIR* func = static_cast<FuncIR*>(expr);
+  auto args = func->all_args();
+  auto str = args[0];
+  auto substr = args[1];
+
+  if (!Match(substr, String(""))) {
+    return false;
+  }
+
+  DCHECK_EQ(ir_graph->dag().ParentsOf(node_id).size(), 1UL);
+  auto parent_id = ir_graph->dag().ParentsOf(node_id)[0];
+  auto parent_node = ir_graph->Get(parent_id);
+
+  PX_RETURN_IF_ERROR(ir_graph->DeleteNode(func->id()));
+
+  // str and substr may be connected to other nodes in the
+  // IR graph. We should only delete those nodes once their
+  // last parent is removed. This should happen with the
+  // preceding FuncIR deletion.
+  if (ir_graph->dag().ParentsOf(str->id()).size() == 0) {
+    PX_RETURN_IF_ERROR(ir_graph->DeleteNode(str->id()));
+  }
+  if (ir_graph->dag().ParentsOf(substr->id()).size() == 0) {
+    PX_RETURN_IF_ERROR(ir_graph->DeleteNode(substr->id()));
+  }
+
+  // Reparent any child nodes of the filter
+  for (int64_t child_id : ir_graph->dag().DependenciesOf(node_id)) {
+    auto child_node = ir_graph->Get(child_id);
+
+    if (!Match(child_node, Operator())) {
+      continue;
+    }
+    auto child_op_node = static_cast<OperatorIR*>(child_node);
+
+    PX_RETURN_IF_ERROR(child_op_node->ReplaceParent(static_cast<OperatorIR*>(ir_node),
+                                                    static_cast<OperatorIR*>(parent_node)));
+  }
+  PX_RETURN_IF_ERROR(ir_graph->DeleteNode(node_id));
+  return true;
+}
+
+StatusOr<bool> PruneUnusedContainsRule::Apply(IRNode* ir_node) {
+  return RemoveMatchingFilter(ir_node);
+}
+
+}  // namespace compiler
+}  // namespace planner
+}  // namespace carnot
+}  // namespace px

src/carnot/planner/compiler/optimizer/prune_unused_contains_rule.h

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <vector>
+
+#include "src/carnot/planner/rules/rules.h"
+
+namespace px {
+namespace carnot {
+namespace planner {
+namespace compiler {
+
+/**
+ * @brief This rule removes contains functions that compare against an empty string.
+ * This results in returning the original dataset and so it is equivalent to no filter
+ * as seen in the example below:
+ *
+ * df = df[px.contains(df.pod, "")]
+ *
+ */
+class PruneUnusedContainsRule : public Rule {
+ public:
+  PruneUnusedContainsRule()
+      : Rule(nullptr, /*use_topo*/ true, /*reverse_topological_execution*/ true) {}
+
+ protected:
+  StatusOr<bool> Apply(IRNode* ir_node) override;
+  StatusOr<bool> RemoveMatchingFilter(IRNode* ir_node);
+};
+
+}  // namespace compiler
+}  // namespace planner
+}  // namespace carnot
+}  // namespace px

src/carnot/planner/compiler/optimizer/prune_unused_contains_rule_test.cc

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2018- The Pixie Authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <string>
+
+#include <gtest/gtest.h>
+
+#include "src/carnot/planner/compiler/analyzer/analyzer.h"
+#include "src/carnot/planner/compiler/optimizer/prune_unused_contains_rule.h"
+#include "src/carnot/planner/compiler/test_utils.h"
+
+namespace px {
+namespace carnot {
+namespace planner {
+namespace compiler {
+
+using table_store::schema::Relation;
+
+using PruneUnusedContainsRuleTest = RulesTest;
+
+TEST_F(PruneUnusedContainsRuleTest, Basic) {
+  MemorySourceIR* mem_src = MakeMemSource(MakeRelation());
+  compiler_state_->relation_map()->emplace("table", MakeRelation());
+
+  auto constant1 = graph->CreateNode<StringIR>(ast, "testing 1").ValueOrDie();
+  auto constant2 = graph->CreateNode<StringIR>(ast, "testing 2").ValueOrDie();
+  auto empty_str = graph->CreateNode<StringIR>(ast, "").ValueOrDie();
+
+  auto filter_func_1 =
+      graph
+          ->CreateNode<FuncIR>(ast, FuncIR::Op{FuncIR::Opcode::non_op, "", "contains"},
+                               std::vector<ExpressionIR*>{constant1, empty_str})
+          .ValueOrDie();
+
+  auto filter_func_2 =
+      graph
+          ->CreateNode<FuncIR>(ast, FuncIR::Op{FuncIR::Opcode::non_op, "", "contains"},
+                               std::vector<ExpressionIR*>{constant2, empty_str})
+          .ValueOrDie();
+
+  auto filter_ir_1 = graph->CreateNode<FilterIR>(ast, mem_src, filter_func_1).ValueOrDie();
+  auto filter_ir_2 = graph->CreateNode<FilterIR>(ast, filter_ir_1, filter_func_2).ValueOrDie();
+
+  auto filter_ir_1_id = filter_ir_1->id();
+  auto filter_ir_2_id = filter_ir_2->id();
+  auto filter_func_ir_1_id = filter_func_1->id();
+  auto filter_func_ir_2_id = filter_func_2->id();
+
+  auto limit = MakeLimit(static_cast<OperatorIR*>(filter_ir_2), 1000);
+
+  ResolveTypesRule type_rule(compiler_state_.get());
+  ASSERT_OK(type_rule.Execute(graph.get()));
+
+  PruneUnusedContainsRule rule;
+  auto result = rule.Execute(graph.get());
+  ASSERT_OK(result);
+  ASSERT_TRUE(result.ConsumeValueOrDie());
+
+  EXPECT_FALSE(graph->HasNode(filter_ir_1_id));
+  EXPECT_FALSE(graph->HasNode(filter_ir_2_id));
+  EXPECT_FALSE(graph->HasNode(filter_func_ir_1_id));
+  EXPECT_FALSE(graph->HasNode(filter_func_ir_2_id));
+  EXPECT_TRUE(graph->HasNode(limit->id()));
+}
+
+TEST_F(PruneUnusedContainsRuleTest, IgnoresNonEmptyStringContains) {
+  MemorySourceIR* mem_src = MakeMemSource(MakeRelation());
+  compiler_state_->relation_map()->emplace("table", MakeRelation());
+
+  auto constant1 = graph->CreateNode<StringIR>(ast, "testing").ValueOrDie();
+  auto constant2 = graph->CreateNode<StringIR>(ast, "existing value").ValueOrDie();
+
+  auto filter_func =
+      graph
+          ->CreateNode<FuncIR>(ast, FuncIR::Op{FuncIR::Opcode::non_op, "", "contains"},
+                               std::vector<ExpressionIR*>{constant1, constant2})
+          .ValueOrDie();
+  auto filter_func_id = filter_func->id();
+
+  auto filter_ir = graph->CreateNode<FilterIR>(ast, mem_src, filter_func).ValueOrDie();
+  auto filter_ir_id = filter_ir->id();
+
+  ResolveTypesRule type_rule(compiler_state_.get());
+  ASSERT_OK(type_rule.Execute(graph.get()));
+
+  PruneUnusedContainsRule rule;
+  auto result = rule.Execute(graph.get());
+  ASSERT_OK(result);
+  ASSERT_FALSE(result.ConsumeValueOrDie());
+
+  EXPECT_TRUE(graph->HasNode(filter_ir_id));
+  EXPECT_TRUE(graph->HasNode(filter_func_id));
+}
+
+}  // namespace compiler
+}  // namespace planner
+}  // namespace carnot
+}  // namespace px

src/carnot/planner/logical_planner_test.cc

@@ -673,6 +673,63 @@ TEST_F(LogicalPlannerTest, filter_pushdown_bug) {
   ASSERT_OK(plan->ToProto());
 }
 
+const char kHttpDataScript[] = R"pxl(
+import px
+
+
+def http_data(start_time: str, source_filter: str, destination_filter: str, num_head: int):
+
+    df = px.DataFrame(table='http_events', start_time=start_time)
+
+    # Add context.
+    df.node = df.ctx['node']
+    df.pid = px.upid_to_pid(df.upid)
+
+    # Filter out entities as specified by the user.
+    df = df[px.contains("source", source_filter)]
+    df = df[px.contains("destination", destination_filter)]
+
+    # Add additional filters below:
+
+    # Restrict number of results.
+    df = df.head(num_head)
+
+    # Order columns.
+
+    return df
+)pxl";
+
+TEST_F(LogicalPlannerTest, VerifyEmptyContainsCallsDoNotSegFaultTest) {
+  auto planner = LogicalPlanner::Create(info_).ConsumeValueOrDie();
+  plannerpb::QueryRequest req;
+  req.set_query_str(kHttpDataScript);
+  auto f = req.add_exec_funcs();
+  f->set_func_name("http_data");
+  f->set_output_table_prefix("http_data");
+  auto start_time = f->add_arg_values();
+  start_time->set_name("start_time");
+  start_time->set_value("-5m");
+
+  auto source_filter = f->add_arg_values();
+  source_filter->set_name("source_filter");
+  source_filter->set_value("");
+
+  auto dest_filter = f->add_arg_values();
+  dest_filter->set_name("destination_filter");
+  dest_filter->set_value("");
+
+  auto num_head = f->add_arg_values();
+  num_head->set_name("num_head");
+  num_head->set_value("1000");
+
+  *req.mutable_logical_planner_state() =
+      testutils::CreateTwoPEMsOneKelvinPlannerState(testutils::kHttpEventsSchema);
+  auto plan_or_s = planner->Plan(req);
+  ASSERT_OK(plan_or_s);
+  auto plan = plan_or_s.ConsumeValueOrDie();
+  EXPECT_OK(plan->ToProto());
+}
+
 TEST_F(LogicalPlannerTest, create_compiler_state_has_endpoint_config) {
   auto state = testutils::CreateTwoPEMsOneKelvinPlannerState(kCheckoutProbeTableSchema);
   auto endpoint_config = state.mutable_otel_endpoint_config();