SubhamSinghal commented on code in PR #27136:
URL: https://github.com/apache/datafusion/pull/27136#discussion_r4501530727
##########
datafusion/physical-plan/src/topk/mod.rs:
##########
@@ -2827,379 +2761,154 @@ impl PartitionedTopKDenseRank {
metrics: &ExecutionPlanMetricsSet,
ranked_schema: Option<SchemaRef>,
) -> Result<Self> {
- assert!(k > 0, "PartitionedTopKDenseRank requires k > 0");
- let reservation =
-
MemoryConsumer::new(format!("PartitionedTopKDenseRank[{partition_id}]"))
- .register(&runtime.memory_pool);
-
- let order_sort_fields = build_sort_fields(&order_expr, schema)?;
- let row_converter = RowConverter::new(order_sort_fields)?;
- let scratch_rows =
- row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW *
batch_size);
-
- let partition_converter = RowConverter::new(partition_sort_fields)?;
- let partition_scratch_rows = partition_converter
- .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size);
-
Ok(Self {
- emit_ranks: ranked_schema.is_some(),
- output_schema: ranked_schema.unwrap_or_else(|| Arc::clone(schema)),
- metrics: TopKMetrics::new(metrics, partition_id),
- reservation,
- expr: order_expr,
- row_converter,
- scratch_rows,
- partition_exprs,
- partition_converter,
- partition_scratch_rows,
- states: HashMap::new(),
- partition_groups: HashMap::new(),
- ob_runs: HashMap::new(),
- store: RecordBatchStore::new(),
- k,
- batch_size,
+ core: PartitionedTopKCore::try_new(
+ "PartitionedTopKDenseRank",
+ partition_id,
+ schema,
+ partition_exprs,
+ partition_sort_fields,
+ order_expr,
+ k,
+ batch_size,
+ runtime,
+ metrics,
+ ranked_schema,
+ )?,
+ groups_bytes: 0,
})
}
- /// Encode PARTITION BY and ORDER BY columns once, demultiplex the
- /// batch's rows by partition key, then per partition bucket the rows
- /// by distinct ob value and merge each bucket into the partition
- /// state as one [`GroupEntry`].
+ /// Encode the partition and ORDER BY columns once each for the whole
batch,
+ /// run every row through its partition's group admission, then register
+ /// the admitted rows in the shared store with one gather.
pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> {
Review Comment:
Addressed in
https://github.com/apache/datafusion/commit/5f96ec422cf781200fb32420e09c6aaa6ed4fed8
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
