Struct TopK 

pub struct TopK {
    schema: SchemaRef,
    metrics: TopKMetrics,
    reservation: MemoryReservation,
    batch_size: usize,
    expr: LexOrdering,
    row_converter: RowConverter,
    scratch_rows: Rows,
    heap: TopKHeap,
    common_sort_prefix_converter: Option<RowConverter>,
    common_sort_prefix: Arc<[PhysicalSortExpr]>,
    filter: Arc<RwLock<TopKDynamicFilters>>,
    pub(crate) finished: bool,
}

Global TopK

Background

“Top K” is a common query optimization used for queries such as “find the top 3 customers by revenue”. The (simplified) SQL for such a query might be:

SELECT customer_id, revenue FROM 'sales.csv' ORDER BY revenue DESC limit 3;

The simple plan would be:

> explain SELECT customer_id, revenue FROM sales ORDER BY revenue DESC limit 3;
+--------------+----------------------------------------+
| plan_type    | plan                                   |
+--------------+----------------------------------------+
| logical_plan | Limit: 3                               |
|              |   Sort: revenue DESC NULLS FIRST       |
|              |     Projection: customer_id, revenue   |
|              |       TableScan: sales                 |
+--------------+----------------------------------------+

While this plan produces the correct answer, it will fully sorts the input before discarding everything other than the top 3 elements.

The same answer can be produced by simply keeping track of the top K=3 elements, reducing the total amount of required buffer memory.

Partial Sort Optimization

This implementation additionally optimizes queries where the input is already partially sorted by a common prefix of the requested ordering. Once the top K heap is full, if subsequent rows are guaranteed to be strictly greater (in sort order) on this prefix than the largest row currently stored, the operator safely terminates early.

Example

For input sorted by (day DESC), but not by timestamp, a query such as:

SELECT day, timestamp FROM sensor ORDER BY day DESC, timestamp DESC LIMIT 10;

can terminate scanning early once sufficient rows from the latest days have been collected, skipping older data.

Structure

This operator tracks the top K items using a TopKHeap.

Fields

schema: SchemaRef

schema of the output (and the input)

metrics: TopKMetrics

Runtime metrics

reservation: MemoryReservation

Reservation

batch_size: usize

The target number of rows for output batches

expr: LexOrdering

sort expressions

row_converter: RowConverter

row converter, for sort keys

scratch_rows: Rows

scratch space for converting rows

heap: TopKHeap

stores the top k values and their sort key values, in order

common_sort_prefix_converter: Option<RowConverter>

row converter, for common keys between the sort keys and the input ordering

common_sort_prefix: Arc<[PhysicalSortExpr]>

Common sort prefix between the input and the sort expressions to allow early exit optimization

filter: Arc<RwLock<TopKDynamicFilters>>

Filter matching the state of the TopK heap used for dynamic filter pushdown

finished: bool

If true, indicates that all rows of subsequent batches are guaranteed to be greater (by byte order, after row conversion) than the top K, which means the top K won’t change and the computation can be finished early.

Implementations

impl TopK

pub fn try_new( partition_id: usize, schema: SchemaRef, common_sort_prefix: Vec<PhysicalSortExpr>, expr: LexOrdering, k: usize, batch_size: usize, runtime: Arc<RuntimeEnv>, metrics: &ExecutionPlanMetricsSet, filter: Arc<RwLock<TopKDynamicFilters>>, ) -> Result<Self>

Create a new TopK that stores the top k values, as defined by the sort expressions in expr.

pub fn insert_batch(&mut self, batch: RecordBatch) -> Result<()>

Insert batch, remembering if any of its values are among the top k seen so far.

fn find_new_topk_items( &mut self, items: impl Iterator<Item = usize>, batch_entry: &mut RecordBatchEntry, ) -> usize

fn update_filter(&mut self) -> Result<()>

Update the filter representation of our TopK heap. For example, given the sort expression ORDER BY a DESC, b ASC LIMIT 3, and the current heap values [(1, 5), (1, 4), (2, 3)], the filter will be updated to:

(a > 1 OR (a = 1 AND b < 5)) AND
(a > 1 OR (a = 1 AND b < 4)) AND
(a > 2 OR (a = 2 AND b < 3))

fn build_filter_expression( sort_exprs: &[PhysicalSortExpr], thresholds: Vec<ScalarValue>, ) -> Result<Option<Arc<dyn PhysicalExpr>>>

Build the filter expression with the given thresholds. This is now called outside of any locks to reduce critical section time.

fn attempt_early_completion(&mut self, batch: &RecordBatch) -> Result<()>

If input ordering shares a common sort prefix with the TopK, and if the TopK’s heap is full, check if the computation can be finished early. This is the case if the last row of the current batch is strictly greater than the max row in the heap, comparing only on the shared prefix columns.

fn compute_common_sort_prefix( &self, batch: &RecordBatch, last_row_idx: usize, scratch: &mut Rows, ) -> Result<()>

pub fn emit(self) -> Result<SendableRecordBatchStream>

Returns the top k results broken into batch_size [RecordBatch]es, consuming the heap

fn size(&self) -> usize

return the size of memory used by this operator, in bytes