comphead commented on code in PR #26067:
URL: https://github.com/apache/datafusion/pull/26067#discussion_r4351060580
##########
datafusion/expr/src/window_state.rs:
##########
@@ -343,12 +348,68 @@ impl PartitionBatchState {
#[derive(Debug, Default, Clone)]
pub struct WindowFrameStateRange {
sort_options: Vec<SortOptions>,
+ peer_comparator: Option<RangePeerComparator>,
+}
+
+/// Reuse comparators to avoid per-row setup work, including logical-null masks
+/// that can take O(n) to build for dictionary arrays.
+#[derive(Clone)]
+struct RangePeerComparator {
+ columns: Vec<ArrayRef>,
+ comparator: Arc<LexicographicalComparator>,
Review Comment:
Nothing in production shares this comparator, so the `Arc` looks like it
exists only for the unit test, which observes reuse through `Arc::downgrade`.
Owning the `LexicographicalComparator` directly and checking reuse through
behavior would keep the type shaped by the runtime need.
##########
datafusion/core/tests/execution/window_nested_range.rs:
##########
@@ -0,0 +1,321 @@
+// 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.
+
+//! Nested RANGE peers through public physical execution APIs.
+//! Construct plans directly to test peer boundaries independently of SQL type
admission.
+
+use std::sync::Arc;
+use std::time::Duration;
+
+use arrow::array::{ArrayRef, AsArray, Int32Array, Int64Array, ListArray,
StructArray};
+use arrow::buffer::{NullBuffer, OffsetBuffer};
+use arrow::compute::concat_batches;
+use arrow::datatypes::{DataType, Field, Int64Type, SchemaRef};
+use arrow::record_batch::RecordBatch;
+use datafusion::datasource::memory::MemorySourceConfig;
+use datafusion::datasource::source::DataSourceExec;
+use datafusion::execution::TaskContext;
+use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
+use datafusion::physical_plan::streaming::{PartitionStream,
StreamingTableExec};
+use datafusion::physical_plan::windows::{
+ BoundedWindowAggExec, WindowAggExec, create_window_expr,
+};
+use datafusion::physical_plan::{
+ ExecutionPlan, ExecutionPlanProperties, InputOrderMode,
SendableRecordBatchStream,
+ collect,
+};
+use datafusion::prelude::SessionContext;
+use datafusion_common::{Result, ScalarValue};
+use datafusion_expr::{
+ WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition,
+};
+use datafusion_functions_aggregate::sum::sum_udaf;
+use datafusion_physical_expr::PhysicalSortExpr;
+use datafusion_physical_expr::expressions::col;
+use datafusion_physical_expr_common::sort_expr::LexOrdering;
+use futures::{FutureExt, StreamExt};
+
+fn nested_batches() -> Result<Vec<RecordBatch>> {
+ // Three nested keys: a NULL container, a nested NULL, and nested 1.
+ // Hidden values under the two NULL containers deliberately differ.
+ let values = [Some(10), Some(99), None, None, None, Some(1)];
+ let values: ArrayRef = Arc::new(Int32Array::from(values.to_vec()));
+ let lists: ArrayRef = Arc::new(ListArray::new(
+ Arc::new(Field::new("item", DataType::Int32, true)),
+ OffsetBuffer::from_lengths([1; 6]),
+ Arc::clone(&values),
+ None,
+ ));
+ let structs: ArrayRef = Arc::new(StructArray::new(
+ vec![Arc::new(Field::new("item", DataType::Int32, true))].into(),
+ vec![values],
+ None,
+ ));
+ let nulls = NullBuffer::from(vec![false, false, true, true, true, true]);
+ let list_field = Arc::new(Field::new("item", lists.data_type().clone(),
true));
+ let keys: [ArrayRef; 3] = [
+ Arc::new(ListArray::new(
+ Arc::clone(&list_field),
+ OffsetBuffer::from_lengths([1; 6]),
+ Arc::clone(&lists),
+ Some(nulls.clone()),
+ )),
+ Arc::new(ListArray::new(
+ Arc::new(Field::new("item", structs.data_type().clone(), true)),
+ OffsetBuffer::from_lengths([1; 6]),
+ structs,
+ Some(nulls.clone()),
+ )),
+ Arc::new(StructArray::new(
+ vec![list_field].into(),
+ vec![lists],
+ Some(nulls),
+ )),
+ ];
+ keys.into_iter()
+ .map(|keys| {
+ // The tie separates the third nested NULL from the preceding
peers.
+ let tie: ArrayRef = Arc::new(Int32Array::from(vec![0, 0, 0, 0, 1,
0]));
+ Ok(RecordBatch::try_from_iter(vec![
+ ("key", keys),
+ ("tie", tie),
+ // In Linear mode, the final row advances another partition's
peers.
+ (
+ "p",
+ Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0, 1])) as
ArrayRef,
+ ),
+ (
+ "value",
+ Arc::new(Int64Array::from(vec![1, 2, 4, 8, 16, 32])) as
ArrayRef,
+ ),
+ ])?)
+ })
+ .collect()
+}
+
+#[tokio::test]
+async fn nested_range_current_row_physical_operators() -> Result<()> {
+ use WindowFrameBound::{CurrentRow, Following, Preceding};
+
+ for batch in nested_batches()? {
+ let schema = batch.schema();
+ let order_by = vec![
+ PhysicalSortExpr::new_default(col("key", &schema)?),
+ PhysicalSortExpr::new_default(col("tie", &schema)?),
+ ];
+ // Repeated peers cross input batch boundaries.
+ let batches = (0..6).map(|row| batch.slice(row, 1)).collect();
+ let source = MemorySourceConfig::try_new(&[batches],
Arc::clone(&schema), None)?
+ .try_with_sort_information(vec![
+ LexOrdering::new(order_by.clone()).unwrap(),
+ ])?;
+ let input: Arc<dyn ExecutionPlan> =
DataSourceExec::from_data_source(source);
+ for (start, end, expected) in [
+ (
+ Preceding(ScalarValue::UInt64(None)),
+ CurrentRow,
+ [3, 3, 15, 15, 31, 63],
+ ),
+ (CurrentRow, CurrentRow, [3, 3, 12, 12, 16, 32]),
+ (
+ CurrentRow,
+ Following(ScalarValue::UInt64(None)),
+ [63, 63, 60, 60, 48, 32],
+ ),
+ ] {
+ for mode in [
+ None,
+ Some(InputOrderMode::Sorted),
+ Some(InputOrderMode::Linear),
+ ] {
+ // The planner uses WindowAggExec for UNBOUNDED FOLLOWING.
+ if mode.is_some() && end.is_unbounded() {
+ continue;
+ }
+ let partition_by = if mode == Some(InputOrderMode::Linear) {
+ vec![col("p", &schema)?]
+ } else {
+ vec![]
+ };
+ let expr = create_window_expr(
+ &WindowFunctionDefinition::AggregateUDF(sum_udaf()),
+ "sum".to_string(),
+ &[col("value", &schema)?],
+ &partition_by,
+ &order_by,
+ Arc::new(WindowFrame::new_bounds(
+ WindowFrameUnits::Range,
+ start.clone(),
+ end.clone(),
+ )),
+ Arc::clone(&schema),
+ false,
+ false,
+ None,
+ )?;
+ let plan: Arc<dyn ExecutionPlan> = if let Some(mode) = &mode {
+ Arc::new(BoundedWindowAggExec::try_new(
+ vec![expr],
+ Arc::clone(&input),
+ mode.clone(),
+ false,
+ )?)
+ } else {
+ Arc::new(WindowAggExec::try_new(
+ vec![expr],
+ Arc::clone(&input),
+ false,
+ )?)
+ };
+ assert!(
+ input
+ .equivalence_properties()
+ .ordering_satisfy_requirement(
+ plan.required_input_ordering()[0]
+ .clone()
+ .unwrap()
+ .into_single(),
+ )?
+ );
+ let output_schema = plan.schema();
+ let output = collect(plan,
SessionContext::new().task_ctx()).await?;
+ let output = concat_batches(&output_schema, &output)?;
+ let actual = output
+ .column(batch.num_columns())
+ .as_primitive::<Int64Type>();
+ let mut expected = expected;
+ if mode == Some(InputOrderMode::Linear) {
+ // The final row is in its own PARTITION BY group.
+ expected[5] = 32;
+ }
+ assert_eq!(
+ actual.iter().collect::<Vec<_>>(),
+ expected.map(Some),
+ "{}, {start:?} to {end:?}, mode={mode:?}",
+ batch.column(0).data_type(),
+ );
+ }
+ }
+ }
+ Ok(())
+}
+
+#[derive(Debug)]
+struct OpenEndedPartition {
+ batch: RecordBatch,
+}
+
+impl PartitionStream for OpenEndedPartition {
+ fn schema(&self) -> &SchemaRef {
+ self.batch.schema_ref()
+ }
+
+ fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
+ let batches = (0..self.batch.num_rows())
+ .map(|row| Ok(self.batch.slice(row, 1)))
+ .collect::<Vec<_>>();
+ Box::pin(RecordBatchStreamAdapter::new(
+ self.batch.schema(),
+ futures::stream::iter(batches).chain(futures::stream::pending()),
+ ))
+ }
+}
+
+#[tokio::test]
+async fn nested_range_emits_completed_peers_before_eof() -> Result<()> {
+ for batch in nested_batches()? {
+ for mode in [InputOrderMode::Sorted, InputOrderMode::Linear] {
+ let schema = batch.schema();
+ let order_by = vec![
+ PhysicalSortExpr::new_default(col("key", &schema)?),
Review Comment:
Every case orders ascending. `is_row_ahead` now depends on the direction
through `is_lt()` with the sort options, so one `DESC` case with the opposite
null placement in this streaming test would pin that down.
--
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]
