comphead commented on code in PR #7461:
URL: https://github.com/apache/datafusion-comet/pull/7461#discussion_r4459282392
##########
spark/src/test/scala/org/apache/comet/exec/CometWindowExecSuite.scala:
##########
@@ -1287,6 +1287,41 @@ class CometWindowExecSuite extends CometTestBase {
}
}
+ test("window: streaming ROW_NUMBER across input batches") {
+ withTempDir { dir =>
+ (0 until 97)
+ .map(i => (if (i < 9) null else s"key-${i / 13}", i, "payload-" * 32))
+ .toDF("k", "i", "payload")
+ .coalesce(1)
+ .write
+ .parquet(dir.toString)
+
spark.read.parquet(dir.toString).createOrReplaceTempView("streaming_row_number")
+ for (batchSize <- Seq(3, 16, 128); ansi <- Seq("false", "true")) {
Review Comment:
This test uses one string key. Most of the new risk is per key type, since
the cross-batch check goes through `make_comparator` and `partition` for
whatever type the key has. Could it move to a SQL fixture under
`sql-tests/windows/` with `-- ConfigMatrix: spark.comet.batchSize=3,16,128`,
with `ROW_NUMBER()` partitioned by int, bigint, decimal, date, timestamp,
string and binary keys, each with NULLs? The float cases from the earlier
review could live there too.
The ANSI half of the loop runs the same plans as far as I can see. Nothing
in the window serde or in `WindowBuilder` reads ANSI for `ROW_NUMBER`, and the
result cast in `window.rs` is built with no eval mode.
##########
native/operators/src/row_number.rs:
##########
@@ -0,0 +1,331 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+
+//! Batch-at-a-time ROW_NUMBER for input sorted by its normalised partition
keys.
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, RecordBatch, UInt64Array};
+use arrow::compute::{partition, take, SortOptions};
+use arrow::datatypes::{DataType, SchemaRef};
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{internal_err, plan_err, Result, ScalarValue,
Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::functions_window::row_number::RowNumber;
+use datafusion::physical_expr::window::StandardWindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr};
+use datafusion::physical_plan::execution_plan::{
+ CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
+};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements,
+ InputOrderMode, PlanProperties, SendableRecordBatchStream,
+};
+use futures::StreamExt;
+
+/// Generates every ROW_NUMBER result in one pass over each input batch.
+///
+/// The unexecuted window supplies the schema, ordering, distribution and
statistics contracts.
+/// Input must use `InputOrderMode::Sorted` so each partition is contiguous.
Its expressions must
+/// all be the built-in ROW_NUMBER over the same partition and sort keys.
+/// The caller must normalise floating-point partition keys before building
the window, just as
+/// it does for ordinary native windows. Other window functions keep their
existing executor.
+#[derive(Debug)]
+pub struct RowNumberExec {
+ window: BoundedWindowAggExec,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl RowNumberExec {
+ pub fn supports(window: &BoundedWindowAggExec) -> bool {
+ if window.input_order_mode != InputOrderMode::Sorted {
+ return false;
+ }
+ let Some(first) = window.window_expr().first() else {
+ return false;
+ };
+ window.window_expr().iter().all(|expr| {
+ let Some(standard) =
expr.as_any().downcast_ref::<StandardWindowExpr>() else {
+ return false;
+ };
+ let Some(udf) = standard
+ .get_standard_func_expr()
+ .as_any()
+ .downcast_ref::<WindowUDFExpr>()
+ else {
+ return false;
+ };
+ udf.fun().inner().downcast_ref::<RowNumber>().is_some()
+ && expr.expressions().is_empty()
+ && expr.partition_by() == first.partition_by()
+ && expr.order_by() == first.order_by()
+ })
+ }
+
+ pub fn try_new(window: BoundedWindowAggExec) -> Result<Self> {
+ if !Self::supports(&window) {
+ return plan_err!(
+ "RowNumberExec requires InputOrderMode::Sorted and only
built-in ROW_NUMBER \
+ expressions with matching keys"
+ );
+ }
+ Ok(Self {
+ window,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+}
+
+impl DisplayAs for RowNumberExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result {
+ write!(f, "CometRowNumberExec: ")?;
+ self.window.fmt_as(t, f)
+ }
+}
+
+impl ExecutionPlan for RowNumberExec {
+ fn name(&self) -> &str {
+ "CometRowNumberExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.window.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ self.window.children()
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ self.window.apply_expressions(f)
+ }
+
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ self.window.required_input_ordering()
+ }
+
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ self.window
+ .input_distribution_requirements()
+ .into_per_child()
+ }
+
+ fn input_distribution_requirements(&self) -> InputDistributionRequirements
{
+ self.window.input_distribution_requirements()
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ options: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let replaced =
Arc::new(self.window.clone()).replace_children(children, options)?;
+ let Some(window) = replaced.downcast_ref::<BoundedWindowAggExec>()
else {
+ return internal_err!("ROW_NUMBER child replacement changed its
window type");
+ };
+ Ok(Arc::new(Self::try_new(window.clone())?))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn with_new_children_and_same_properties(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let input = self
+ .window
+ .input()
+ .execute(partition, Arc::clone(&context))?;
+ let partition_by =
self.window.window_expr()[0].partition_by().to_vec();
+ let output_count = self.window.window_expr().len();
+ let schema = self.schema();
+ let state = RowNumberState {
+ key: vec![],
+ rank: 0,
+ reservation: MemoryConsumer::new("ROW_NUMBER partition key")
+ .register(context.memory_pool()),
+ };
+ let metrics = BaselineMetrics::new(&self.metrics, partition);
+ let stream = futures::stream::try_unfold(
+ (input, state, metrics),
+ move |(mut input, mut state, metrics)| {
+ let partition_by = partition_by.clone();
+ let schema = Arc::clone(&schema);
+ async move {
+ let Some(batch) = input.next().await else {
+ return Ok(None);
+ };
+ let batch = batch?;
+ let output = {
+ let _timer = metrics.elapsed_compute().timer();
+ state.evaluate(batch, &partition_by, schema,
output_count)?
+ };
+ (&output).record_output(&metrics);
+ Ok(Some((output, (input, state, metrics))))
+ }
+ },
+ );
+ Ok(Box::pin(RecordBatchStreamAdapter::new(
+ self.schema(),
+ stream,
+ )))
+ }
+
+ fn metrics(&self) -> Option<MetricsSet> {
+ Some(self.metrics.clone_inner())
+ }
+
+ fn child_stats_requests(&self, partition: Option<usize>) ->
Vec<ChildStats> {
+ self.window.child_stats_requests(partition)
+ }
+
+ fn statistics_from_inputs(
+ &self,
+ input_stats: &[Arc<Statistics>],
+ args: &StatisticsArgs,
+ ) -> Result<Arc<Statistics>> {
+ self.window.statistics_from_inputs(input_stats, args)
+ }
+}
+
+struct RowNumberState {
+ key: Vec<ArrayRef>,
+ rank: u64,
+ reservation: MemoryReservation,
+}
+
+impl RowNumberState {
+ fn evaluate(
+ &mut self,
+ batch: RecordBatch,
+ partition_by: &[Arc<dyn PhysicalExpr>],
+ schema: SchemaRef,
+ output_count: usize,
+ ) -> Result<RecordBatch> {
+ let rows = batch.num_rows();
+ let mut ranks = Vec::with_capacity(rows);
+ if rows != 0 {
+ let keys = partition_by
+ .iter()
+ .map(|expr| expr.evaluate(&batch)?.into_array(rows))
+ .collect::<Result<Vec<_>>>()?;
+ let mut continues = keys.len() == self.key.len();
+ for (key, previous) in keys.iter().zip(&self.key) {
+ let compare = arrow::array::make_comparator(
+ previous.as_ref(),
+ key.as_ref(),
+ SortOptions::default(),
+ )?;
+ continues &= compare(0, 0).is_eq();
+ }
+ let ranges = if keys.is_empty() {
+ std::iter::once(0..rows).collect::<Vec<_>>()
+ } else {
+ partition(&keys)?.ranges()
+ };
+ for range in &ranges {
+ let start = if range.start == 0 && continues {
+ self.rank
+ } else {
+ 0
+ };
+ ranks.extend((1..=range.len() as u64).map(|offset|
start.wrapping_add(offset)));
+ }
+ self.rank = *ranks.last().unwrap();
+ if !keys.is_empty() && (!continues || ranges.len() != 1) {
+ // The old boundary key is no longer needed. Do not keep both
copies alive
+ // when a wide key changes. A continued partition needs no new
key copy.
+ self.key = vec![];
+ self.reservation.free();
+ let index = UInt64Array::from(vec![(rows - 1) as u64]);
+ let key = keys
+ .iter()
+ .map(|array| {
+ if matches!(
+ array.data_type(),
+ DataType::Dictionary(_, _) | DataType::Utf8View |
DataType::BinaryView
+ ) {
+ // A view or dictionary take retains its source
buffers. Rebuild
+ // the logical value so unrelated bytes do not
stay alive.
+ ScalarValue::try_from_array(array.as_ref(), rows -
1)?.to_array()
+ } else {
+ take(array.as_ref(), &index,
None).map_err(Into::into)
+ }
+ })
+ .collect::<Result<Vec<_>>>()?;
+ // Nested views and dictionaries can still share backing
buffers after take.
+ // Charge the buffers retained, rather than assuming a one-row
key is small.
+ let size = key
+ .iter()
+ .map(|array| array.get_array_memory_size())
+ .sum::<usize>()
+ + key.capacity() * std::mem::size_of::<ArrayRef>();
+ self.reservation.try_resize(size)?;
Review Comment:
With `free()` here and `try_resize` at line 315, every batch that ends a
partition hands the key's bytes back to the pool and asks for them again. On
`fair_unified` (the default) and `greedy_unified`, `shrink` and `try_grow` each
call Spark over JNI (`fair_pool.rs:218-279`, `unified_pool.rs:99-120`), and the
acquire runs in `block_in_place`. At small partition sizes that is two JNI
calls per batch. The benchmark uses DataFusion's default pool, so it would not
show this, and I have not measured it. Clearing `self.key` without calling
`free()` still drops the old copy first, and the resize then reports only the
change in size, which is zero for fixed-width keys.
Also, the key already exists when it is charged and this operator cannot
spill. `memory_management.md` (line 485) describes `grow` for exactly that
case. With `try_resize`, a refusal fails a task that `BoundedWindowAggExec`,
which reserves nothing, would have run. Would `resize` fit better here?
`refused_boundary_reservations_are_released` would then go away.
##########
native/operators/src/row_number.rs:
##########
@@ -0,0 +1,331 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+
+//! Batch-at-a-time ROW_NUMBER for input sorted by its normalised partition
keys.
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, RecordBatch, UInt64Array};
+use arrow::compute::{partition, take, SortOptions};
+use arrow::datatypes::{DataType, SchemaRef};
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{internal_err, plan_err, Result, ScalarValue,
Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::functions_window::row_number::RowNumber;
+use datafusion::physical_expr::window::StandardWindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr};
+use datafusion::physical_plan::execution_plan::{
+ CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
+};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements,
+ InputOrderMode, PlanProperties, SendableRecordBatchStream,
+};
+use futures::StreamExt;
+
+/// Generates every ROW_NUMBER result in one pass over each input batch.
+///
+/// The unexecuted window supplies the schema, ordering, distribution and
statistics contracts.
+/// Input must use `InputOrderMode::Sorted` so each partition is contiguous.
Its expressions must
+/// all be the built-in ROW_NUMBER over the same partition and sort keys.
+/// The caller must normalise floating-point partition keys before building
the window, just as
+/// it does for ordinary native windows. Other window functions keep their
existing executor.
+#[derive(Debug)]
+pub struct RowNumberExec {
+ window: BoundedWindowAggExec,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl RowNumberExec {
+ pub fn supports(window: &BoundedWindowAggExec) -> bool {
+ if window.input_order_mode != InputOrderMode::Sorted {
+ return false;
+ }
+ let Some(first) = window.window_expr().first() else {
+ return false;
+ };
+ window.window_expr().iter().all(|expr| {
+ let Some(standard) =
expr.as_any().downcast_ref::<StandardWindowExpr>() else {
+ return false;
+ };
+ let Some(udf) = standard
+ .get_standard_func_expr()
+ .as_any()
+ .downcast_ref::<WindowUDFExpr>()
+ else {
+ return false;
+ };
+ udf.fun().inner().downcast_ref::<RowNumber>().is_some()
+ && expr.expressions().is_empty()
+ && expr.partition_by() == first.partition_by()
+ && expr.order_by() == first.order_by()
+ })
+ }
+
+ pub fn try_new(window: BoundedWindowAggExec) -> Result<Self> {
+ if !Self::supports(&window) {
+ return plan_err!(
+ "RowNumberExec requires InputOrderMode::Sorted and only
built-in ROW_NUMBER \
+ expressions with matching keys"
+ );
+ }
+ Ok(Self {
+ window,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+}
+
+impl DisplayAs for RowNumberExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result {
+ write!(f, "CometRowNumberExec: ")?;
+ self.window.fmt_as(t, f)
+ }
+}
+
+impl ExecutionPlan for RowNumberExec {
+ fn name(&self) -> &str {
+ "CometRowNumberExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.window.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ self.window.children()
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ self.window.apply_expressions(f)
+ }
+
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ self.window.required_input_ordering()
+ }
+
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ self.window
+ .input_distribution_requirements()
+ .into_per_child()
+ }
+
+ fn input_distribution_requirements(&self) -> InputDistributionRequirements
{
+ self.window.input_distribution_requirements()
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ options: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let replaced =
Arc::new(self.window.clone()).replace_children(children, options)?;
+ let Some(window) = replaced.downcast_ref::<BoundedWindowAggExec>()
else {
+ return internal_err!("ROW_NUMBER child replacement changed its
window type");
+ };
+ Ok(Arc::new(Self::try_new(window.clone())?))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn with_new_children_and_same_properties(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let input = self
+ .window
+ .input()
+ .execute(partition, Arc::clone(&context))?;
+ let partition_by =
self.window.window_expr()[0].partition_by().to_vec();
Review Comment:
On main the native window fails with `Expects PARTITION BY expression to be
ordered` when its input does not declare the order of the partition keys.
`BoundedWindowAggExec::execute` reaches `get_search_algo`, which calls
`partition_by_sort_keys()` (DataFusion 55.2
`bounded_window_agg_exec.rs:268-293` and `windows/mod.rs:354-372`).
`RowNumberExec::execute` reads `self.window.input()` directly and skips that
check.
I tried it in a scratch crate with this file, over a `MemorySourceConfig`
with no sort information and keys `[1, 2, 1, 2]` in one batch.
`BoundedWindowAggExec` returns that error, and `RowNumberExec` returns `1, 1,
1, 1` where Spark gives `1, 2` per key. When the rows really are sorted and
only the declaration is missing, as in #7223, the new path would give the right
answer. When the order Spark reports is false (#6898), a partition that arrives
in two runs is numbered from 1 twice and the query succeeds. #7256 also relies
on this check (its `CometWindowExec` comment says the native window rejects an
input whose order it cannot see).
Could `execute` call `self.window.partition_by_sort_keys()?` first, with a
test over a source without sort information next to
`interleaved_partitions_require_the_generic_executor`?
##########
native/operators/src/row_number.rs:
##########
@@ -0,0 +1,331 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you 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.
+
+//! Batch-at-a-time ROW_NUMBER for input sorted by its normalised partition
keys.
+
+use std::fmt::Formatter;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, RecordBatch, UInt64Array};
+use arrow::compute::{partition, take, SortOptions};
+use arrow::datatypes::{DataType, SchemaRef};
+use datafusion::common::tree_node::TreeNodeRecursion;
+use datafusion::common::{internal_err, plan_err, Result, ScalarValue,
Statistics};
+use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
+use datafusion::execution::TaskContext;
+use datafusion::functions_window::row_number::RowNumber;
+use datafusion::physical_expr::window::StandardWindowExpr;
+use datafusion::physical_expr::{OrderingRequirements, PhysicalExpr};
+use datafusion::physical_plan::execution_plan::{
+ CardinalityEffect, ChildrenPropertiesMode, ReplaceChildrenOptions,
+};
+use datafusion::physical_plan::metrics::{
+ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
+};
+use datafusion::physical_plan::statistics::{ChildStats, StatisticsArgs};
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr};
+use datafusion::physical_plan::{
+ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan,
InputDistributionRequirements,
+ InputOrderMode, PlanProperties, SendableRecordBatchStream,
+};
+use futures::StreamExt;
+
+/// Generates every ROW_NUMBER result in one pass over each input batch.
+///
+/// The unexecuted window supplies the schema, ordering, distribution and
statistics contracts.
+/// Input must use `InputOrderMode::Sorted` so each partition is contiguous.
Its expressions must
+/// all be the built-in ROW_NUMBER over the same partition and sort keys.
+/// The caller must normalise floating-point partition keys before building
the window, just as
+/// it does for ordinary native windows. Other window functions keep their
existing executor.
+#[derive(Debug)]
+pub struct RowNumberExec {
+ window: BoundedWindowAggExec,
+ metrics: ExecutionPlanMetricsSet,
+}
+
+impl RowNumberExec {
+ pub fn supports(window: &BoundedWindowAggExec) -> bool {
+ if window.input_order_mode != InputOrderMode::Sorted {
+ return false;
+ }
+ let Some(first) = window.window_expr().first() else {
+ return false;
+ };
+ window.window_expr().iter().all(|expr| {
+ let Some(standard) =
expr.as_any().downcast_ref::<StandardWindowExpr>() else {
+ return false;
+ };
+ let Some(udf) = standard
+ .get_standard_func_expr()
+ .as_any()
+ .downcast_ref::<WindowUDFExpr>()
+ else {
+ return false;
+ };
+ udf.fun().inner().downcast_ref::<RowNumber>().is_some()
+ && expr.expressions().is_empty()
+ && expr.partition_by() == first.partition_by()
+ && expr.order_by() == first.order_by()
+ })
+ }
+
+ pub fn try_new(window: BoundedWindowAggExec) -> Result<Self> {
+ if !Self::supports(&window) {
+ return plan_err!(
+ "RowNumberExec requires InputOrderMode::Sorted and only
built-in ROW_NUMBER \
+ expressions with matching keys"
+ );
+ }
+ Ok(Self {
+ window,
+ metrics: ExecutionPlanMetricsSet::new(),
+ })
+ }
+}
+
+impl DisplayAs for RowNumberExec {
+ fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) ->
std::fmt::Result {
+ write!(f, "CometRowNumberExec: ")?;
+ self.window.fmt_as(t, f)
+ }
+}
+
+impl ExecutionPlan for RowNumberExec {
+ fn name(&self) -> &str {
+ "CometRowNumberExec"
+ }
+
+ fn properties(&self) -> &Arc<PlanProperties> {
+ self.window.properties()
+ }
+
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ self.window.children()
+ }
+
+ fn apply_expressions(
+ &self,
+ f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ self.window.apply_expressions(f)
+ }
+
+ fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
+ self.window.required_input_ordering()
+ }
+
+ fn required_input_distribution(&self) -> Vec<Distribution> {
+ self.window
+ .input_distribution_requirements()
+ .into_per_child()
+ }
+
+ fn input_distribution_requirements(&self) -> InputDistributionRequirements
{
+ self.window.input_distribution_requirements()
+ }
+
+ fn maintains_input_order(&self) -> Vec<bool> {
+ vec![true]
+ }
+
+ fn cardinality_effect(&self) -> CardinalityEffect {
+ CardinalityEffect::Equal
+ }
+
+ fn replace_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ options: ReplaceChildrenOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let replaced =
Arc::new(self.window.clone()).replace_children(children, options)?;
+ let Some(window) = replaced.downcast_ref::<BoundedWindowAggExec>()
else {
+ return internal_err!("ROW_NUMBER child replacement changed its
window type");
+ };
+ Ok(Arc::new(Self::try_new(window.clone())?))
+ }
+
+ fn with_new_children(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
+ )
+ }
+
+ fn with_new_children_and_same_properties(
+ self: Arc<Self>,
+ children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ self.replace_children(
+ children,
+ ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep),
+ )
+ }
+
+ fn execute(
+ &self,
+ partition: usize,
+ context: Arc<TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ let input = self
+ .window
+ .input()
+ .execute(partition, Arc::clone(&context))?;
+ let partition_by =
self.window.window_expr()[0].partition_by().to_vec();
+ let output_count = self.window.window_expr().len();
+ let schema = self.schema();
+ let state = RowNumberState {
+ key: vec![],
+ rank: 0,
+ reservation: MemoryConsumer::new("ROW_NUMBER partition key")
+ .register(context.memory_pool()),
+ };
+ let metrics = BaselineMetrics::new(&self.metrics, partition);
+ let stream = futures::stream::try_unfold(
+ (input, state, metrics),
+ move |(mut input, mut state, metrics)| {
+ let partition_by = partition_by.clone();
Review Comment:
`partition_by.clone()` allocates a new `Vec` for every batch. If
`partition_by`, `schema` and `output_count` live in `RowNumberState`,
`evaluate` can borrow them and this closure needs no per-batch clones.
--
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]
