sunchao commented on code in PR #26067: URL: https://github.com/apache/datafusion/pull/26067#discussion_r4349939947
########## datafusion/core/tests/execution/window_nested_range.rs: ########## @@ -0,0 +1,251 @@ +// 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. Review Comment: Moved the tests into tests/execution/window_nested_range.rs and registered the module in the existing core_integration target. The separate integration-test binary is gone. Updated in [`9910027`](https://github.com/apache/datafusion/pull/26067/commits/99100274dd139574854e40393474caac5c97e16b). ########## datafusion/core/tests/window_nested_range.rs: ########## @@ -0,0 +1,251 @@ +// 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, 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), + ( + "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 bounded in [false, true] { + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("value", &schema)?], + &[], + &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 bounded { + Arc::new(BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(&input), + InputOrderMode::Sorted, Review Comment: Added Linear coverage for frames ending at CURRENT ROW for all three nested key types in the existing tests. A newer row in a second partition must now release completed peers in the first before EOF, with the ordering requirement checked. The fix uses Arrow comparison and reuses its Boolean result within each partition evaluation, avoiding repeated comparator setup without persistent state. The genuine PartiallySorted SQL control passes already; that stream path does not populate most_recent_row, so it does not reach this scalar comparison with a row bound. Updated in [`9910027`](https://github.com/apache/datafusion/pull/26067/commits/99100274dd139574854e40393474caac5c97e16b). -- 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]
