1919//! values without materializing normalized columns.
2020
2121use crate :: float_semantics:: {
22- has_float_leaf , normalize_comparison_operand, normalize_nested_floats, spark_equality ,
23- NormalizeNestedFloats ,
22+ is_nested_with_float_leaf , normalize_comparison_operand, normalize_nested_floats,
23+ spark_equality , NormalizeNestedFloats ,
2424} ;
2525use arrow:: array:: { Array , ArrayRef , AsArray , BooleanArray } ;
2626use arrow:: buffer:: { BooleanBuffer , NullBuffer } ;
@@ -29,17 +29,13 @@ use arrow::datatypes::{DataType, Schema};
2929use arrow:: record_batch:: RecordBatch ;
3030use datafusion:: common:: { internal_err, DFSchema , Result , ScalarValue } ;
3131use datafusion:: logical_expr:: { ColumnarValue , Operator } ;
32- use datafusion:: physical_expr:: expressions:: { in_list, BinaryExpr , InListExpr } ;
32+ use datafusion:: physical_expr:: expressions:: { in_list, BinaryExpr , Column , InListExpr , Literal } ;
3333use datafusion:: physical_expr:: PhysicalExpr ;
3434use datafusion:: physical_expr_common:: physical_expr:: is_volatile;
3535use std:: fmt:: { Display , Formatter } ;
3636use std:: hash:: { Hash , Hasher } ;
3737use std:: sync:: Arc ;
3838
39- fn needs_spark_equality ( dt : & DataType ) -> bool {
40- dt. is_nested ( ) && has_float_leaf ( dt)
41- }
42-
4339struct Operand {
4440 array : ArrayRef ,
4541 scalar : bool ,
@@ -226,9 +222,13 @@ impl PhysicalExpr for NestedPredicate {
226222pub enum FloatOperands {
227223 /// Normalize them, so that the comparison follows Spark's SQL ordering.
228224 Normalize ,
229- /// Leave them as they are. Only a scan's pushed-down data filters use this: Parquet pruning
230- /// recognizes a column compared with a literal but not a normalized column, and Spark's
231- /// Filter above the scan evaluates each of those filters again.
225+ /// Leave a Float32 or Float64 column compared with a literal as it is, and normalize every
226+ /// other operand. Only a scan's pushed-down data filters use this: Parquet pruning recognizes
227+ /// a column compared with a literal but not a normalized column. With row-level pushdown the
228+ /// reader also evaluates the filters on each row, and a row it drops never reaches Spark's
229+ /// Filter above the scan, so any other shape, which pruning cannot use anyway, is normalized.
230+ /// A computed operand such as `-d` can hold a NaN with the sign bit set, which a raw
231+ /// comparison sorts below every other value.
232232 Raw ,
233233}
234234
@@ -260,7 +260,7 @@ pub fn spark_comparison(
260260 let ( Ok ( left_type) , Ok ( _) ) = ( left. data_type ( schema) , right. data_type ( schema) ) else {
261261 return Ok ( Arc :: new ( BinaryExpr :: new ( left, op, right) ) ) ;
262262 } ;
263- if matches ! ( op, Eq | NotEq ) && needs_spark_equality ( & left_type) {
263+ if matches ! ( op, Eq | NotEq ) && is_nested_with_float_leaf ( & left_type) {
264264 validate_types ( & left, std:: slice:: from_ref ( & right) , schema) ?;
265265 return Ok ( Arc :: new ( NestedPredicate {
266266 value : left,
@@ -269,16 +269,33 @@ pub fn spark_comparison(
269269 membership : false ,
270270 } ) ) ;
271271 }
272- let ( left, right) = match float_operands {
273- FloatOperands :: Normalize => (
272+ let raw = float_operands == FloatOperands :: Raw ;
273+ let ( left, right) = if raw && is_float_column ( & left, schema) && is_literal ( & right) {
274+ ( left, normalize_comparison_operand ( right, schema) ?)
275+ } else if raw && is_literal ( & left) && is_float_column ( & right, schema) {
276+ ( normalize_comparison_operand ( left, schema) ?, right)
277+ } else {
278+ (
274279 normalize_comparison_operand ( left, schema) ?,
275280 normalize_comparison_operand ( right, schema) ?,
276- ) ,
277- FloatOperands :: Raw => ( left, right) ,
281+ )
278282 } ;
279283 Ok ( Arc :: new ( BinaryExpr :: new ( left, op, right) ) )
280284}
281285
286+ /// Whether `expr` is a Float32 or Float64 column, the operand that Parquet pruning reads.
287+ fn is_float_column ( expr : & Arc < dyn PhysicalExpr > , schema : & Schema ) -> bool {
288+ expr. downcast_ref :: < Column > ( ) . is_some ( )
289+ && matches ! (
290+ expr. data_type( schema) ,
291+ Ok ( DataType :: Float32 | DataType :: Float64 )
292+ )
293+ }
294+
295+ fn is_literal ( expr : & Arc < dyn PhysicalExpr > ) -> bool {
296+ expr. downcast_ref :: < Literal > ( ) . is_some ( )
297+ }
298+
282299fn validate_types (
283300 value : & Arc < dyn PhysicalExpr > ,
284301 candidates : & [ Arc < dyn PhysicalExpr > ] ,
@@ -301,7 +318,7 @@ pub fn spark_in_list(
301318 negated : bool ,
302319 schema : & Schema ,
303320) -> Result < Arc < dyn PhysicalExpr > > {
304- if !needs_spark_equality ( & value. data_type ( schema) ?) || candidates. is_empty ( ) {
321+ if !is_nested_with_float_leaf ( & value. data_type ( schema) ?) || candidates. is_empty ( ) {
305322 return in_list ( value, candidates, & negated, schema) ;
306323 }
307324 validate_types ( & value, & candidates, schema) ?;
@@ -948,7 +965,8 @@ mod tests {
948965 }
949966
950967 /// A literal operand is normalized while planning, so the comparison stays `column op
951- /// literal`. With `FloatOperands::Raw`, as for a scan's data filters, nothing is wrapped.
968+ /// literal`. With `FloatOperands::Raw`, as for a scan's data filters, only the column of that
969+ /// shape is left unwrapped.
952970 #[ test]
953971 fn float_operand_shapes ( ) -> Result < ( ) > {
954972 use crate :: float_semantics:: NormalizeNaNAndZero ;
@@ -974,6 +992,8 @@ mod tests {
974992 let folded = binary. right ( ) . downcast_ref :: < Literal > ( ) . unwrap ( ) ;
975993 assert ! ( matches!( folded. value( ) , ScalarValue :: Float64 ( Some ( v) ) if v. to_bits( ) == 0 ) ) ;
976994
995+ // With `FloatOperands::Raw`, a float column compared with a literal keeps the column, on
996+ // either side, and the literal is still normalized.
977997 let expr = spark_comparison (
978998 Arc :: new ( Column :: new ( "a" , 0 ) ) ,
979999 Operator :: Lt ,
@@ -983,10 +1003,39 @@ mod tests {
9831003 ) ?;
9841004 let binary = expr. downcast_ref :: < BinaryExpr > ( ) . unwrap ( ) ;
9851005 assert ! ( binary. left( ) . downcast_ref:: <Column >( ) . is_some( ) ) ;
986- assert ! ( Arc :: ptr_eq( binary. right( ) , & negative_zero) ) ;
1006+ let folded = binary. right ( ) . downcast_ref :: < Literal > ( ) . unwrap ( ) ;
1007+ assert ! ( matches!( folded. value( ) , ScalarValue :: Float64 ( Some ( v) ) if v. to_bits( ) == 0 ) ) ;
1008+ let expr = spark_comparison (
1009+ Arc :: clone ( & negative_zero) ,
1010+ Operator :: Lt ,
1011+ Arc :: new ( Column :: new ( "a" , 0 ) ) ,
1012+ & schema,
1013+ FloatOperands :: Raw ,
1014+ ) ?;
1015+ let binary = expr. downcast_ref :: < BinaryExpr > ( ) . unwrap ( ) ;
1016+ assert ! ( binary. right( ) . downcast_ref:: <Column >( ) . is_some( ) ) ;
9871017
988- // Other operators are not comparisons and keep their operands.
1018+ // Any other shape is normalized even with `FloatOperands::Raw`, because a reader with
1019+ // row-level pushdown drops the rows it rejects.
9891020 let ( a, b) = columns ( ) ;
1021+ let expr = spark_comparison (
1022+ Arc :: clone ( & a) ,
1023+ Operator :: Lt ,
1024+ Arc :: clone ( & b) ,
1025+ & schema,
1026+ FloatOperands :: Raw ,
1027+ ) ?;
1028+ let binary = expr. downcast_ref :: < BinaryExpr > ( ) . unwrap ( ) ;
1029+ assert ! ( binary
1030+ . left( )
1031+ . downcast_ref:: <NormalizeNaNAndZero >( )
1032+ . is_some( ) ) ;
1033+ assert ! ( binary
1034+ . right( )
1035+ . downcast_ref:: <NormalizeNaNAndZero >( )
1036+ . is_some( ) ) ;
1037+
1038+ // Other operators are not comparisons and keep their operands.
9901039 let expr = spark_comparison (
9911040 Arc :: clone ( & a) ,
9921041 Operator :: Plus ,
0 commit comments