Add --use_zero_copy_filter flag, which allows PEMs to perform zero-copy filters but with signficant batching. This trades off memory allocations for incrased compute time (decreases data locality and batch efficiency)

Signed-off-by: Dom Del Nano <ddelnano@gmail.com>

Author: ddelnano

src/carnot/exec/filter_node.cc

@@ -39,6 +39,8 @@
 #include "src/shared/types/types.h"
 #include "src/shared/types/typespb/wrapper/types_pb_wrapper.h"
 
+DEFINE_bool(use_zero_copy_filter, false, "Uses a non optimized, but zero-copy filter that can run on PEMs. If this setting is disabled, this assumes that filters must run on Kelvin");
+
 namespace px {
 namespace carnot {
 namespace exec {
@@ -132,42 +134,66 @@ Status PredicateCopyValues<types::STRING>(const types::BoolValueColumnWrapper& p
 Status FilterNode::ConsumeNextImpl(ExecState* exec_state, const RowBatch& rb, size_t) {
   // Current implementation does not merge across row batches, we should
   // consider this for cases where the filter has really low selectivity.
-  /* PX_ASSIGN_OR_RETURN(auto pred_col, evaluator_->EvaluateSingleExpression( */
-  /*                                        exec_state, rb, *plan_node_->expression())); */
-
-  /* // Verify that the type of the column is boolean. */
-  /* DCHECK_EQ(pred_col->data_type(), types::BOOLEAN) << "Predicate expression must be a boolean"; */
-
-  /* const types::BoolValueColumnWrapper& pred_col_wrapper = */
-  /*     *static_cast<types::BoolValueColumnWrapper*>(pred_col.get()); */
-  /* size_t num_pred = pred_col_wrapper.Size(); */
-
-  /* DCHECK_EQ(static_cast<size_t>(rb.num_rows()), num_pred); */
-
-  // Find out how many of them returned true;
-  size_t num_output_records = rb.num_rows();
-  /* for (size_t i = 0; i < num_pred; ++i) { */
-  /*   if (pred_col_wrapper[i].val) { */
-  /*     ++num_output_records; */
-  /*   } */
-  /* } */
-
-  RowBatch output_rb(*output_descriptor_, num_output_records);
-  DCHECK_EQ(output_descriptor_->size(), plan_node_->selected_cols().size());
-
-  for (const auto& [output_col_idx, input_col_idx] : Enumerate(plan_node_->selected_cols())) {
-    auto input_col = rb.ColumnAt(input_col_idx);
-    PX_RETURN_IF_ERROR(output_rb.AddColumn(input_col));
-    /* auto col_type = output_descriptor_->type(output_col_idx); */
-/* #define TYPE_CASE(_dt_) \ */
-/*   PX_RETURN_IF_ERROR(PredicateCopyValues<_dt_>(pred_col_wrapper, input_col.get(), &output_rb)); */
-/*     PX_SWITCH_FOREACH_DATATYPE(col_type, TYPE_CASE); */
-/* #undef TYPE_CASE */
-  }
+  if (FLAGS_use_zero_copy_filter) {
+    PX_ASSIGN_OR_RETURN(auto pred_col, evaluator_->EvaluateSingleExpression(
+                                           exec_state, rb, *plan_node_->expression()));
+
+    // Verify that the type of the column is boolean.
+    DCHECK_EQ(pred_col->data_type(), types::BOOLEAN) << "Predicate expression must be a boolean";
+
+    const types::BoolValueColumnWrapper& pred_col_wrapper =
+        *static_cast<types::BoolValueColumnWrapper*>(pred_col.get());
+    size_t num_pred = pred_col_wrapper.Size();
+
+    DCHECK_EQ(static_cast<size_t>(rb.num_rows()), num_pred);
+
+    std::vector<RowBatch> output_rbs;
+    int64_t segment_start = -1;
+    for (size_t idx = 0; idx < num_pred; ++idx) {
+      if (pred_col_wrapper[idx].val) {
+        if (segment_start == -1) {
+          segment_start = idx;
+        }
+      } else {
+        if (segment_start != -1) {
+          output_rbs.emplace_back(rb.UnsafeSlice(segment_start, idx - segment_start));
+          segment_start = -1;
+        }
+      }
+    }
+
+    if (output_rbs.empty() && segment_start != -1) {
+      output_rbs.emplace_back(rb.UnsafeSlice(segment_start, num_pred - segment_start));
+    } else if (output_rbs.empty()) {
+      // No rows matched the filter.
+      PX_ASSIGN_OR_RETURN(auto empty_rb,
+                          RowBatch::WithZeroRows(*output_descriptor_, rb.eow(), rb.eos()));
+      RowBatch& empty_rb_ref = *empty_rb;
+      output_rbs.push_back(empty_rb_ref);
+    }
 
-  output_rb.set_eow(rb.eow());
-  output_rb.set_eos(rb.eos());
-  PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, output_rb));
+    auto& back = output_rbs.back();
+    back.set_eow(rb.eow());
+    back.set_eos(rb.eos());
+
+    for (auto it = output_rbs.begin(); it != output_rbs.end(); ++it) {
+      auto& output_rb = *it;
+      PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, output_rb));
+    }
+  } else {
+    size_t num_output_records = rb.num_rows();
+    RowBatch output_rb(*output_descriptor_, num_output_records);
+    DCHECK_EQ(output_descriptor_->size(), plan_node_->selected_cols().size());
+
+    for (const auto& [output_col_idx, input_col_idx] : Enumerate(plan_node_->selected_cols())) {
+      auto input_col = rb.ColumnAt(input_col_idx);
+      PX_RETURN_IF_ERROR(output_rb.AddColumn(input_col));
+    }
+
+    output_rb.set_eow(rb.eow());
+    output_rb.set_eos(rb.eos());
+    PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, output_rb));
+  }
   return Status::OK();
 }
 

src/carnot/exec/filter_node.h

@@ -32,6 +32,8 @@
 #include "src/common/base/status.h"
 #include "src/table_store/table_store.h"
 
+DECLARE_bool(use_zero_copy_filter);
+
 namespace px {
 namespace carnot {
 namespace exec {

src/table_store/schema/row_batch.cc

@@ -265,6 +265,15 @@ StatusOr<std::unique_ptr<RowBatch>> RowBatch::Slice(int64_t offset, int64_t leng
   return output_rb;
 }
 
+RowBatch RowBatch::UnsafeSlice(int64_t offset, int64_t length) const {
+  auto output_rb = RowBatch(desc(), length);
+  for (int64_t input_col_idx = 0; input_col_idx < num_columns(); ++input_col_idx) {
+    auto col = ColumnAt(input_col_idx);
+    PX_UNUSED(output_rb.AddColumn(col->Slice(offset, length)));
+  }
+  return output_rb;
+}
+
 }  // namespace schema
 }  // namespace table_store
 }  // namespace px

src/table_store/schema/row_batch.h

@@ -75,6 +75,7 @@ class RowBatch {
    * @return StatusOr<std::unique_ptr<RowBatch>>
    */
   StatusOr<std::unique_ptr<RowBatch>> Slice(int64_t offset, int64_t length) const;
+  RowBatch UnsafeSlice(int64_t offset, int64_t length) const;
 
   /**
    * Adds the given column to the row batch, given that it correctly fits the schema.