viirya commented on code in PR #26067:
URL: https://github.com/apache/datafusion/pull/26067#discussion_r4319630746
##########
datafusion/expr/src/window_state.rs:
##########
@@ -343,12 +347,66 @@ impl PartitionBatchState {
#[derive(Debug, Default, Clone)]
pub struct WindowFrameStateRange {
sort_options: Vec<SortOptions>,
+ peer_comparator: Option<RangePeerComparator>,
+}
+
+/// Cache comparators because dictionary logical-null masks can scan entire
arrays.
+#[derive(Clone)]
+struct RangePeerComparator {
+ columns: Vec<ArrayRef>,
+ comparator: Arc<LexicographicalComparator>,
+}
+
+// The cached Arrow arrays and comparison closures are immutable. A comparison
+// that unwinds cannot leave this cache partially updated.
+impl std::panic::UnwindSafe for RangePeerComparator {}
Review Comment:
Which public type does the downstream `catch_unwind` reproducer need to stay
`UnwindSafe`? `WindowAggState` holds an `ArrayRef` (`out_col`), and
`WindowFrameContext::Groups` holds `ScalarValue`s, so by the auto-trait rules
neither looks `UnwindSafe` already. That would leave `WindowFrameStateRange` on
its own (constructible through `Default`) as the only type affected. If that's
the case, these manual impls may not be buying much.
##########
datafusion/expr/src/window_state.rs:
##########
@@ -70,7 +71,10 @@ impl WindowAggState {
match self.window_frame_ctx.as_mut() {
// Rows have no state do nothing
Some(WindowFrameContext::Rows(_)) => {}
- Some(WindowFrameContext::Range { .. }) => {}
+ Some(WindowFrameContext::Range { state, .. }) => {
+ // The comparator refers to the old, unpruned ORDER BY arrays.
+ state.peer_comparator = None;
Review Comment:
Since the cache is keyed on `Arc::ptr_eq`, the sliced arrays after pruning
would already miss the cache. So this is only about releasing the old arrays
early, not correctness. Maybe say that in the comment.
##########
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:
Only `InputOrderMode::Sorted` is covered. Adding a `Linear` case with a
PARTITION BY column would catch the `is_row_ahead` failure mentioned in the
summary.
##########
datafusion/expr/src/window_state.rs:
##########
@@ -343,12 +347,66 @@ impl PartitionBatchState {
#[derive(Debug, Default, Clone)]
pub struct WindowFrameStateRange {
sort_options: Vec<SortOptions>,
+ peer_comparator: Option<RangePeerComparator>,
+}
+
+/// Cache comparators because dictionary logical-null masks can scan entire
arrays.
Review Comment:
Could this be simpler? GROUPS frames already find peers with `ScalarValue`
equality (`check_equality` / `search_in_slice`), and `ScalarValue::eq` goes
through `nested_eq` for lists, structs and maps. Reusing that would avoid the
cache, the `prune_state` hook, the manual auto-trait impls and the custom
`Debug`, at the same per-row cost as the existing non-nested path. If there is
a semantic reason it doesn't work here (dictionary nulls?), a comment saying so
would help. If the cache stays, I think the general reason is that building a
comparator per row is costly, and O(n) for dictionaries. The comment currently
gives only the dictionary case.
##########
datafusion/expr/src/window_state.rs:
##########
@@ -440,6 +498,28 @@ impl WindowFrameStateRange {
delta: Option<&ScalarValue>,
length: usize,
) -> Result<usize> {
+ if delta.is_none()
+ && range_columns
+ .iter()
+ .any(|column| column.data_type().is_nested())
+ {
+ // CURRENT ROW needs peer equality. Scalar ordering does not
support
+ // all nested keys; Arrow's sort comparator also handles
dictionary nulls.
+ let mut boundary = if SIDE {
+ last_range.start
+ } else {
+ last_range.end.max(idx)
Review Comment:
The `.max(idx)` is load-bearing here, unlike in the ordered search below. An
equality scan stops at the first non-peer, so starting before `idx` would
return immediately. Worth a short comment.
##########
datafusion/expr/src/window_state.rs:
##########
@@ -440,6 +498,28 @@ impl WindowFrameStateRange {
delta: Option<&ScalarValue>,
length: usize,
) -> Result<usize> {
+ if delta.is_none()
+ && range_columns
+ .iter()
+ .any(|column| column.data_type().is_nested())
+ {
+ // CURRENT ROW needs peer equality. Scalar ordering does not
support
+ // all nested keys; Arrow's sort comparator also handles
dictionary nulls.
+ let mut boundary = if SIDE {
+ last_range.start
+ } else {
+ last_range.end.max(idx)
+ };
+ let comparator = self.peer_comparator(range_columns)?;
+ while boundary < length {
+ let is_peer = comparator.compare(boundary, idx).is_eq();
+ if if SIDE { is_peer } else { !is_peer } {
Review Comment:
nit: `if is_peer == SIDE { break; }` is easier to read than the nested `if
if`.
##########
datafusion/expr/src/window_state.rs:
##########
@@ -440,6 +498,28 @@ impl WindowFrameStateRange {
delta: Option<&ScalarValue>,
length: usize,
) -> Result<usize> {
+ if delta.is_none()
+ && range_columns
+ .iter()
+ .any(|column| column.data_type().is_nested())
+ {
+ // CURRENT ROW needs peer equality. Scalar ordering does not
support
+ // all nested keys; Arrow's sort comparator also handles
dictionary nulls.
Review Comment:
"Arrow's sort comparator also handles dictionary nulls" reads as if
dictionary keys take this path, but only nested types (including
dictionary-of-nested) do. Could this say what problem it refers to?
##########
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.
Review Comment:
This adds a new integration test binary, which means another link step in
CI. Could these tests live under `core_integration` (e.g. a module in
`tests/physical_optimizer` or `tests/execution`) instead?
##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -1325,11 +1325,8 @@ fn extract_window_frame_target_type(col_type: &DataType)
-> Option<DataType> {
/// run over an ORDER BY column of `col_type` even though the type has no
/// arithmetic for finite offsets.
///
-/// Such a frame only compares rows to find peers, so the type must compare
-/// the same way in the RANGE peer check (`ScalarValue::partial_cmp`) as in
-/// the sort that produced the input order. That holds for durations and
-/// intervals; it does not for structs and maps, whose `ScalarValue`
-/// comparison differs from the sorter's, so they stay unsupported.
+/// This fallback admits durations and intervals, including encoded values.
+/// Structs and maps remain outside its supported types.
Review Comment:
The removed text explained why these types are rejected: the peer check went
through `compare_rows`. For nested keys that is no longer true, so for example
`ORDER BY d, l` would now be computed correctly but is still rejected. I agree
with keeping admission unchanged in this PR, but the comment should say these
restrictions are now conservative and could be relaxed, ideally linking a
follow-up issue, rather than just restating the rule. (Same for `is_list_type`
below.)
##########
datafusion/expr/src/window_state.rs:
##########
@@ -729,6 +809,73 @@ mod tests {
use arrow::array::{Float64Array, UInt64Array};
use arrow::datatypes::{Field, Schema};
+ #[test]
+ fn range_peer_comparator_reuses_arrays_and_invalidates_after_pruning() ->
Result<()> {
Review Comment:
This test asserts implementation details: cache identity through
`Weak::upgrade`, and the exact point where memory is released. That will break
under harmless refactors. The result correctness after pruning (the last
assertion) seems like the part worth keeping.
##########
datafusion/sqllogictest/test_files/window.slt:
##########
@@ -7207,23 +7206,21 @@ select count(*) over (order by l) from (values ([1]),
([2])) t(l)
1
2
-# but a list whose elements are nested (struct, list) plans and then fails at
-# execution: the peer comparison in the RANGE frame goes through
-# `ScalarValue::partial_cmp`, which cannot compare nested list elements, while
-# the sort that produced the input uses `make_comparator` and can. Documented
-# here so the fix to the scalar comparison can flip these expectations.
-query error DataFusion error: Internal error: Uncomparable values:
List\(\[\{c0: 1\}\]\), List\(\[\{c0: 1\}\]\)
+# Lists containing structs or lists also form RANGE peer groups.
Review Comment:
Could we add the plain-list case that this PR fixes (see the review
summary), e.g. `order by l nulls first` over `[NULL]`, `[1]`, `[NULL]`? It
returns 3, 3, 3 on main and 2, 2, 3 here.
--
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]
