[refactor](topn) Refactor topn filter push down

be/src/olap/column_predicate.h

@@ -158,7 +158,7 @@ struct PredicateTypeTraits {
  } \
  }
 
-class ColumnPredicate {
+class ColumnPredicate : public std::enable_shared_from_this<ColumnPredicate> {
 public:
  explicit ColumnPredicate(uint32_t column_id, PrimitiveType primitive_type,
  bool opposite = false)

be/src/olap/rowset/segment_v2/segment.cpp

@@ -275,39 +275,6 @@ Status Segment::new_iterator(SchemaSPtr schema, const StorageReadOptions& read_o
  }
  }
 
- if (!read_options.topn_filter_source_node_ids.empty()) {
- auto* query_ctx = read_options.runtime_state->get_query_ctx();
- for (int id : read_options.topn_filter_source_node_ids) {
- auto runtime_predicate = query_ctx->get_runtime_predicate(id).get_predicate(
- read_options.topn_filter_target_node_id);
-
- AndBlockColumnPredicate and_predicate;
- and_predicate.add_column_predicate(
- SingleColumnBlockPredicate::create_unique(runtime_predicate));
- std::shared_ptr<ColumnReader> reader;
- Status st = get_column_reader(
- read_options.tablet_schema->column(runtime_predicate->column_id()), &reader,
- read_options.stats);
- if (st.is<ErrorCode::NOT_FOUND>()) {
- continue;
- }
- RETURN_IF_ERROR(st);
- DCHECK(reader != nullptr);
- if (can_apply_predicate_safely(runtime_predicate->column_id(), *schema,
- read_options.target_cast_type_for_variants,
- read_options)) {
- bool matched = true;
- RETURN_IF_ERROR(reader->match_condition(&and_predicate, &matched));
- if (!matched) {
- // any condition not satisfied, return.
- *iter = std::make_unique<EmptySegmentIterator>(*schema);
- read_options.stats->filtered_segment_number++;
- return Status::OK();
- }
- }
- }
- }
-
  {
  SCOPED_RAW_TIMER(&read_options.stats->segment_load_index_timer_ns);
  RETURN_IF_ERROR(load_index(read_options.stats));

be/src/olap/rowset/segment_v2/segment.h

@@ -180,6 +180,7 @@ class Segment : public std::enable_shared_from_this<Segment>, public MetadataAdd
  const std::map<std::string, vectorized::DataTypePtr>& target_cast_type_for_variants,
  const StorageReadOptions& read_options) {
  const doris::Field* col = schema.column(cid);
+ DCHECK(col != nullptr) << "Column not found in schema for cid=" << cid;
  vectorized::DataTypePtr storage_column_type =
  get_data_type_of(col->get_desc(), read_options);
  if (storage_column_type == nullptr || col->type() != FieldType::OLAP_FIELD_TYPE_VARIANT ||

be/src/olap/rowset/segment_v2/segment_iterator.cpp

@@ -946,31 +946,6 @@ Status SegmentIterator::_get_row_ranges_from_conditions(RowRanges* condition_row
  RowRanges::ranges_intersection(*condition_row_ranges, zone_map_row_ranges,
  condition_row_ranges);
 
- if (!_opts.topn_filter_source_node_ids.empty()) {
- auto* query_ctx = _opts.runtime_state->get_query_ctx();
- for (int id : _opts.topn_filter_source_node_ids) {
- std::shared_ptr<doris::ColumnPredicate> runtime_predicate =
- query_ctx->get_runtime_predicate(id).get_predicate(
- _opts.topn_filter_target_node_id);
- if (_segment->can_apply_predicate_safely(runtime_predicate->column_id(), *_schema,
- _opts.target_cast_type_for_variants,
- _opts)) {
- AndBlockColumnPredicate and_predicate;
- and_predicate.add_column_predicate(
- SingleColumnBlockPredicate::create_unique(runtime_predicate));
-
- RowRanges column_rp_row_ranges = RowRanges::create_single(num_rows());
- RETURN_IF_ERROR(_column_iterators[runtime_predicate->column_id()]
- ->get_row_ranges_by_zone_map(&and_predicate, nullptr,
- &column_rp_row_ranges));
-
- // intersect different columns's row ranges to get final row ranges by zone map
- RowRanges::ranges_intersection(zone_map_row_ranges, column_rp_row_ranges,
- &zone_map_row_ranges);
- }
- }
- }
-
  size_t pre_size2 = condition_row_ranges->count();
  RowRanges::ranges_intersection(*condition_row_ranges, zone_map_row_ranges,
  condition_row_ranges);
@@ -1688,26 +1663,6 @@ Status SegmentIterator::_vec_init_lazy_materialization() {
  }
  }
 
- // add runtime predicate to _col_predicates
- // should NOT add for order by key,
- // since key is already sorted and topn_next only need first N rows from each segment,
- // but runtime predicate will filter some rows and read more than N rows.
- // should add add for order by none-key column, since none-key column is not sorted and
- // all rows should be read, so runtime predicate will reduce rows for topn node
- if (!_opts.topn_filter_source_node_ids.empty() &&
- (_opts.read_orderby_key_columns == nullptr || _opts.read_orderby_key_columns->empty())) {
- for (int id : _opts.topn_filter_source_node_ids) {
- auto& runtime_predicate =
- _opts.runtime_state->get_query_ctx()->get_runtime_predicate(id);
- _col_predicates.push_back(
- runtime_predicate.get_predicate(_opts.topn_filter_target_node_id));
- VLOG_DEBUG << fmt::format(
- "After appending topn filter to col_predicates, "
- "col_predicates size: {}, col_predicate: {}",
- _col_predicates.size(), _col_predicates.back()->debug_string());
- }
- }
-
  // Step1: extract columns that can be lazy materialization
  if (!_col_predicates.empty() || !del_cond_id_set.empty()) {
  std::set<ColumnId> short_cir_pred_col_id_set; // using set for distinct cid

be/src/olap/shared_predicate.h

@@ -54,8 +54,10 @@ class SharedPredicate final : public ColumnPredicate {
  ColumnPredicate::debug_string(), _nested ? _nested->debug_string() : "null");
  return fmt::to_string(debug_string_buffer);
  }
+ void set_column_id(uint32_t column_id) { _column_id = column_id; }
  std::shared_ptr<ColumnPredicate> clone(uint32_t column_id) const override {
- return SharedPredicate::create_shared(*this, column_id);
+ // All scanner thread should share the same SharedPredicate object.
+ return std::const_pointer_cast<ColumnPredicate>(shared_from_this());
  }
 
  PredicateType type() const override {

be/src/pipeline/exec/olap_scan_operator.h

@@ -79,6 +79,7 @@ class OlapScanLocalState final : public ScanLocalState<OlapScanLocalState> {
  PushDownType& pdt) override;
 
  PushDownType _should_push_down_bloom_filter() override { return PushDownType::ACCEPTABLE; }
+ PushDownType _should_push_down_topn_filter() override { return PushDownType::ACCEPTABLE; }
 
  PushDownType _should_push_down_bitmap_filter() override { return PushDownType::ACCEPTABLE; }
 

be/src/pipeline/exec/scan_operator.cpp

@@ -325,6 +325,19 @@ Status ScanLocalState<Derived>::_normalize_predicate(vectorized::VExprContext* c
  _parent->cast<typename Derived::Parent>()._slot_id_to_slot_desc[slot->slot_id()];
  return _is_predicate_acting_on_slot(slot, range);
  };
+ auto topn_predicate_checker = [&](const vectorized::VExprSPtrs& children,
+ SlotDescriptor** slot_desc, ColumnValueRangeType** range) {
+ if (children.empty() || children[0]->node_type() != TExprNodeType::SLOT_REF) {
+ // not a slot ref(column)
+ return false;
+ }
+ std::shared_ptr<vectorized::VSlotRef> slot =
+ std::dynamic_pointer_cast<vectorized::VSlotRef>(children[0]);
+ CHECK(slot != nullptr);
+ *slot_desc =
+ _parent->cast<typename Derived::Parent>()._slot_id_to_slot_desc[slot->slot_id()];
+ return _is_predicate_acting_on_slot(slot, range);
+ };
 
  if (expr_root != nullptr) {
  if (is_leaf(expr_root)) {
@@ -353,7 +366,8 @@ Status ScanLocalState<Derived>::_normalize_predicate(vectorized::VExprContext* c
  vectorized::VExpr::expr_without_cast(child));
  }
  if (in_predicate_checker(expr_root->children(), &slot, &range) ||
- eq_predicate_checker(expr_root->children(), &slot, &range)) {
+ eq_predicate_checker(expr_root->children(), &slot, &range) ||
+ topn_predicate_checker(expr_root->children(), &slot, &range)) {
  Status status = Status::OK();
  std::visit(
  [&](auto& value_range) {
@@ -403,7 +417,14 @@ Status ScanLocalState<Derived>::_normalize_predicate(vectorized::VExprContext* c
  context, slot,
  _slot_id_to_predicates[slot->id()], &pdt),
  status);
-
+ RETURN_IF_PUSH_DOWN(_normalize_bloom_filter(
+ context, slot,
+ _slot_id_to_predicates[slot->id()], &pdt),
+ status);
+ RETURN_IF_PUSH_DOWN(_normalize_topn_filter(
+ context, slot,
+ _slot_id_to_predicates[slot->id()], &pdt),
+ status);
  if (state()->enable_function_pushdown()) {
  RETURN_IF_PUSH_DOWN(
  _normalize_function_filters(context, slot, &pdt), status);
@@ -465,6 +486,26 @@ Status ScanLocalState<Derived>::_normalize_bloom_filter(
  return Status::OK();
 }
 
+template <typename Derived>
+Status ScanLocalState<Derived>::_normalize_topn_filter(
+ vectorized::VExprContext* expr_ctx, SlotDescriptor* slot,
+ std::vector<std::shared_ptr<ColumnPredicate>>& predicates, PushDownType* pdt) {
+ auto expr = expr_ctx->root()->is_rf_wrapper() ? expr_ctx->root()->get_impl() : expr_ctx->root();
+ if (expr->is_topn_filter()) {
+ PushDownType temp_pdt = _should_push_down_topn_filter();
+ if (temp_pdt != PushDownType::UNACCEPTABLE) {
+ auto& p = _parent->cast<typename Derived::Parent>();
+ auto& pred = _state->get_query_ctx()->get_runtime_predicate(
+ assert_cast<vectorized::VTopNPred*>(expr.get())->source_node_id());
+ if (_push_down_topn(pred)) {
+ predicates.emplace_back(pred.get_predicate(p.node_id()));
+ *pdt = temp_pdt;
+ }
+ }
+ }
+ return Status::OK();
+}
+
 template <typename Derived>
 Status ScanLocalState<Derived>::_normalize_bitmap_filter(
  vectorized::VExprContext* expr_ctx, SlotDescriptor* slot,
@@ -1211,6 +1252,18 @@ Status ScanLocalState<Derived>::_get_topn_filters(RuntimeState* state) {
  RETURN_IF_ERROR(conjunct->open(state));
  _conjuncts.emplace_back(conjunct);
  }
+ for (auto id : get_topn_filter_source_node_ids(state, true)) {
+ const auto& pred = state->get_query_ctx()->get_runtime_predicate(id);
+ vectorized::VExprSPtr topn_pred;
+ RETURN_IF_ERROR(vectorized::VTopNPred::create_vtopn_pred(pred.get_texpr(p.node_id()), id,
+ topn_pred));
+
+ vectorized::VExprContextSPtr conjunct = vectorized::VExprContext::create_shared(topn_pred);
+ RETURN_IF_ERROR(conjunct->prepare(
+ state, _parent->cast<typename Derived::Parent>().row_descriptor()));
+ RETURN_IF_ERROR(conjunct->open(state));
+ _conjuncts.emplace_back(conjunct);
+ }
  return Status::OK();
 }
 
@@ -1335,8 +1388,9 @@ Status ScanOperatorX<LocalStateType>::prepare(RuntimeState* state) {
  continue;
  }
 
- state->get_query_ctx()->get_runtime_predicate(id).init_target(node_id(),
- _slot_id_to_slot_desc);
+ RETURN_IF_ERROR(state->get_query_ctx()->get_runtime_predicate(id).init_target(
+ node_id(), _slot_id_to_slot_desc,
+ OperatorX<LocalStateType>::intermediate_row_desc()));
  }
 
  RETURN_IF_CANCELLED(state);

be/src/pipeline/exec/scan_operator.h

@@ -217,6 +217,7 @@ class ScanLocalState : public ScanLocalStateBase {
  virtual bool _push_down_topn(const vectorized::RuntimePredicate& predicate) { return false; }
  virtual bool _is_key_column(const std::string& col_name) { return false; }
  virtual PushDownType _should_push_down_bloom_filter() { return PushDownType::UNACCEPTABLE; }
+ virtual PushDownType _should_push_down_topn_filter() { return PushDownType::UNACCEPTABLE; }
  virtual PushDownType _should_push_down_bitmap_filter() { return PushDownType::UNACCEPTABLE; }
  virtual PushDownType _should_push_down_is_null_predicate() {
  return PushDownType::UNACCEPTABLE;
@@ -254,6 +255,9 @@ class ScanLocalState : public ScanLocalStateBase {
  Status _normalize_bloom_filter(vectorized::VExprContext* expr_ctx, SlotDescriptor* slot,
  std::vector<std::shared_ptr<ColumnPredicate>>& predicates,
  PushDownType* pdt);
+ Status _normalize_topn_filter(vectorized::VExprContext* expr_ctx, SlotDescriptor* slot,
+ std::vector<std::shared_ptr<ColumnPredicate>>& predicates,
+ PushDownType* pdt);
 
  Status _normalize_bitmap_filter(vectorized::VExprContext* expr_ctx, SlotDescriptor* slot,
  std::vector<std::shared_ptr<ColumnPredicate>>& predicates,

be/src/runtime/runtime_predicate.cpp

@@ -55,16 +55,31 @@ RuntimePredicate::RuntimePredicate(const TTopnFilterDesc& desc)
  : create_comparison_predicate0<PredicateType::GE>;
 }
 
-void RuntimePredicate::init_target(
- int32_t target_node_id, phmap::flat_hash_map<int, SlotDescriptor*> slot_id_to_slot_desc) {
+Status RuntimePredicate::init_target(
+ int32_t target_node_id, phmap::flat_hash_map<int, SlotDescriptor*> slot_id_to_slot_desc,
+ const doris::RowDescriptor& desc) {
  std::unique_lock<std::shared_mutex> wlock(_rwlock);
  check_target_node_id(target_node_id);
  if (target_is_slot(target_node_id)) {
  _contexts[target_node_id].col_name =
  slot_id_to_slot_desc[get_texpr(target_node_id).nodes[0].slot_ref.slot_id]
  ->col_name();
+ auto slot_id = get_texpr(target_node_id).nodes[0].slot_ref.slot_id;
+ auto column_id = desc.get_column_id(slot_id);
+ if (column_id < 0) {
+ return Status::Error<ErrorCode::INTERNAL_ERROR>(
+ "RuntimePredicate has invalid slot id: {}, name: {}, desc: {}, slot_desc: {}",
+ slot_id,
+ slot_id_to_slot_desc[get_texpr(target_node_id).nodes[0].slot_ref.slot_id]
+ ->col_name(),
+ desc.debug_string(),
+ slot_id_to_slot_desc[get_texpr(target_node_id).nodes[0].slot_ref.slot_id]
+ ->debug_string());
+ }
+ _contexts[target_node_id].predicate = SharedPredicate::create_shared(column_id);
  }
  _detected_target = true;
+ return Status::OK();
 }
 
 StringRef RuntimePredicate::_get_string_ref(const Field& field, const PrimitiveType type) {

be/src/runtime/runtime_predicate.h

@@ -44,8 +44,9 @@ class RuntimePredicate {
 public:
  RuntimePredicate(const TTopnFilterDesc& desc);
 
- void init_target(int32_t target_node_id,
- phmap::flat_hash_map<int, SlotDescriptor*> slot_id_to_slot_desc);
+ Status init_target(int32_t target_node_id,
+ phmap::flat_hash_map<int, SlotDescriptor*> slot_id_to_slot_desc,
+ const doris::RowDescriptor& desc);
 
  bool enable() const {
  // when sort node and scan node are not in the same fragment, predicate will be disabled
@@ -66,9 +67,10 @@ class RuntimePredicate {
  }
  RETURN_IF_ERROR(tablet_schema->have_column(_contexts[target_node_id].col_name));
  _contexts[target_node_id].tablet_schema = tablet_schema;
- int64_t index = DORIS_TRY(_contexts[target_node_id].get_field_index())
- _contexts[target_node_id]
- .predicate = SharedPredicate::create_shared(index);
+ int64_t index = DORIS_TRY(_contexts[target_node_id].get_field_index());
+ DCHECK(_contexts[target_node_id].predicate != nullptr);
+ assert_cast<SharedPredicate*>(_contexts[target_node_id].predicate.get())
+ ->set_column_id(cast_set<uint32_t>(index));
  return Status::OK();
  }
 
@@ -130,6 +132,7 @@ class RuntimePredicate {
  struct TargetContext {
  TExpr expr;
  std::string col_name;
+ // TODO(gabriel): remove this
  TabletSchemaSPtr tablet_schema;
  std::shared_ptr<ColumnPredicate> predicate;
 

be/src/vec/exprs/vexpr.h

@@ -209,6 +209,7 @@ class VExpr {
  return std::ranges::any_of(_children.begin(), _children.end(),
  [](VExprSPtr child) { return child->is_rf_wrapper(); });
  }
+ virtual bool is_topn_filter() const { return false; }
 
  virtual void do_judge_selectivity(uint64_t filter_rows, uint64_t input_rows) {
  for (auto child : _children) {

be/src/vec/exprs/vtopn_pred.h

@@ -45,6 +45,7 @@ class VTopNPred : public VExpr {
  _source_node_id(source_node_id),
  _expr_name(fmt::format("VTopNPred(source_node_id={})", _source_node_id)),
  _target_ctx(std::move(target_ctx)) {}
+ bool is_topn_filter() const override { return true; }
 
  static Status create_vtopn_pred(const TExpr& target_expr, int source_node_id,
  vectorized::VExprSPtr& expr) {
@@ -63,6 +64,8 @@ class VTopNPred : public VExpr {
  return Status::OK();
  }
 
+ int source_node_id() const { return _source_node_id; }
+
  Status prepare(RuntimeState* state, const RowDescriptor& desc, VExprContext* context) override {
  _predicate = &state->get_query_ctx()->get_runtime_predicate(_source_node_id);
  RETURN_IF_ERROR_OR_PREPARED(VExpr::prepare(state, desc, context));