kosiew commented on code in PR #26020: URL: https://github.com/apache/datafusion/pull/26020#discussion_r4346029783
########## datafusion/physical-expr/src/expressions/normalize_float_zero.rs: ########## @@ -0,0 +1,230 @@ +// 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. + +//! Canonical floating-point zeros in window partition keys. + +use std::hash::Hash; +use std::sync::Arc; + +use arrow::datatypes::{DataType, FieldRef, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::ColumnarValue; +use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::sort_properties::ExprProperties; + +use crate::PhysicalExpr; + +/// Replaces negative floating-point zeros with positive zeros, including nested values. +/// +/// # Public Only for Internal Use +/// +/// Used by window planning and physical-plan serialization. Normalizing the partition +/// expression makes sorting, partition boundaries, and cross-batch keys agree without +/// changing the original input columns. +#[doc(hidden)] +#[derive(Debug, Eq)] +pub struct NormalizeFloatZeroExpr { + arg: Arc<dyn PhysicalExpr>, +} + +// Manually implement PartialEq and Hash to work around rust-lang/rust#78808. +impl PartialEq for NormalizeFloatZeroExpr { + fn eq(&self, other: &Self) -> bool { + self.arg.eq(&other.arg) + } +} + +impl Hash for NormalizeFloatZeroExpr { + fn hash<H: std::hash::Hasher>(&self, state: &mut H) { + self.arg.hash(state); + } +} + +impl NormalizeFloatZeroExpr { + /// Create a normalized window partition key. + pub fn new(arg: Arc<dyn PhysicalExpr>) -> Self { + Self { arg } + } +} + +impl std::fmt::Display for NormalizeFloatZeroExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "normalize_float_zero({})", self.arg) + } +} + +impl PhysicalExpr for NormalizeFloatZeroExpr { + fn data_type(&self, input_schema: &Schema) -> Result<DataType> { + self.arg.data_type(input_schema) + } + + fn nullable(&self, input_schema: &Schema) -> Result<bool> { + self.arg.nullable(input_schema) + } + + fn return_field(&self, input_schema: &Schema) -> Result<FieldRef> { + self.arg.return_field(input_schema) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> { + Ok(match self.arg.evaluate(batch)? { + ColumnarValue::Array(array) => { + ColumnarValue::Array(normalize_float_zero(&array)) + } + ColumnarValue::Scalar(scalar) if scalar.data_type().is_floating() => { + ColumnarValue::Scalar(normalize_float_zero_scalar(scalar)) + } + ColumnarValue::Scalar(scalar) => { + // Nested scalar values need the recursive array normalization too. + let array = normalize_float_zero(&scalar.to_array()?); + ColumnarValue::Scalar(ScalarValue::try_from_array(&array, 0)?) + } + }) + } + + fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> { + vec![&self.arg] + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn PhysicalExpr>>, + ) -> Result<Arc<dyn PhysicalExpr>> { + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } + + fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> { + let child = &children[0]; + if child.range.data_type().is_floating() { + Ok(ExprProperties { + sort_properties: child.sort_properties, + range: Interval::try_new( + normalize_float_zero_scalar(child.range.lower().clone()), + normalize_float_zero_scalar(child.range.upper().clone()), + )?, + preserves_lex_ordering: true, + // Merging signed-zero groups can invalidate suffix sort keys. + strictly_order_preserving: false, + }) + } else { + // Nested values are not monotone: [-0, 2] < [+0, 1], but [0, 2] > [0, 1]. Review Comment: Could we add a focused ordering-property regression using raw list keys `[-0.0, 2]` before `[+0.0, 1]`? Normalization reverses their ordering, so the test should verify that the normalized key is not considered already sorted. The existing one-element-list test does not cover this inversion. ########## datafusion/proto/tests/cases/plans/windows.rs: ########## @@ -306,3 +306,82 @@ fn roundtrip_lead_with_default_value() -> Result<()> { true, )?)) } + +/// Raw [a, b] ordering is insufficient after signed-zero groups merge. Verify +/// planning adds the required ordering and that it survives serialization. +#[tokio::test] +async fn roundtrip_signed_zero_window() -> Result<()> { + use datafusion::arrow::array::{Float64Array, Int64Array, RecordBatch, UInt64Array}; + use datafusion::datasource::MemTable; + use datafusion::physical_plan::collect; + use datafusion::prelude::{SessionConfig, SessionContext, col}; + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + }; + + let ctx = + SessionContext::new_with_config(SessionConfig::new().with_target_partitions(1)); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Float64Array::from(vec![-0.0, 0.0])), + Arc::new(Int64Array::from(vec![2, 1])), + ], + )?; + let table = MemTable::try_new(schema, vec![vec![batch]])?.with_sort_order(vec![ Review Comment: Could we also test a presorted single-float-key input with `count(*) OVER (PARTITION BY a)` and no window `ORDER BY`? Please verify that signed zeros share the correct counts and that no additional `SortExec` is inserted, protecting reuse of safe prefix ordering. -- 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]
