Core: Fix Z-order byte encoding for floating point values

core/src/main/java/org/apache/iceberg/util/ZOrderByteUtils.java

@@ -102,7 +102,7 @@ public static ByteBuffer doubleToOrderedBytes(double val, ByteBuffer reuse) {
   public static ByteBuffer floatingPointOrderedBytes(double val, ByteBuffer reuse) {
     ByteBuffer bytes = ByteBuffers.reuse(reuse, PRIMITIVE_BUFFER_SIZE);
     long lval = Double.doubleToLongBits(val);
-    lval ^= ((lval >> (Integer.SIZE - 1)) | Long.MIN_VALUE);
+    lval ^= ((lval >> (Long.SIZE - 1)) | Long.MIN_VALUE);
     bytes.putLong(lval);
     return bytes;
   }

core/src/test/java/org/apache/iceberg/util/TestZOrderByteUtil.java

@@ -24,7 +24,10 @@
 import java.nio.charset.CharsetEncoder;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.List;
 import java.util.Random;
+import java.util.function.Function;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.primitives.UnsignedBytes;
 import org.apache.iceberg.types.Types;
 import org.junit.jupiter.api.Test;
@@ -379,6 +382,130 @@ public void testByteTruncateOrFill() {
     }
   }
 
+  /**
+   * The random float/double ordering tests above draw from {@code nextFloat()}/{@code
+   * nextDouble()}, so the compared values essentially never agree on their high bits. Ordering is
+   * then decided by the high bits alone, which hides any corruption of the low bits. These cases
+   * pin the ordering for values that differ only in low mantissa bits, where the low bytes are what
+   * decides the comparison.
+   */
+  @Test
+  public void testFloatOrderingForConsecutiveMantissaValues() {
+    int baseBits = Float.floatToIntBits(1.0f);
+    List<Float> values = Lists.newArrayList();
+    for (int i = 0; i < 64; i++) {
+      values.add(Float.intBitsToFloat(baseBits + i));
+    }
+
+    assertOrderPreserved(values, TestZOrderByteUtil::encodeFloat);
+  }
+
+  @Test
+  public void testDoubleOrderingForValuesDifferingInLowMantissaBits() {
+    List<Double> values = Lists.newArrayList();
+    for (int i = 0; i < 64; i++) {
+      values.add(1.0d + (i * 0x1.0p-30));
+    }
+
+    assertOrderPreserved(values, TestZOrderByteUtil::encodeDouble);
+  }
+
+  @Test
+  public void testNegativeDoubleOrderingForValuesDifferingInLowMantissaBits() {
+    List<Double> values = Lists.newArrayList();
+    for (int i = 63; i >= 0; i--) {
+      values.add(-1.0d - (i * 0x1.0p-30));
+    }
+
+    assertOrderPreserved(values, TestZOrderByteUtil::encodeDouble);
+  }
+
+  /** Boundary pairs, including the ones that straddle zero and the extremes of the range. */
+  @Test
+  public void testDoubleOrderingForBoundaryPairs() {
+    double[][] ascendingPairs = {
+      {-921614.125d, -921614.0625d},
+      {-1.6001329423771755E213d, -1.600132804916327E213d},
+      {5.716890676284865E-207d, 5.7168911255697246E-207d},
+      {-Double.MIN_VALUE, 0.0d},
+      {0.0d, Double.MIN_VALUE},
+      {-Double.MAX_VALUE, Double.MAX_VALUE},
+      {-1.0d, 1.0d},
+    };
+
+    for (double[] pair : ascendingPairs) {
+      assertOrderPreserved(Lists.newArrayList(pair[0], pair[1]), TestZOrderByteUtil::encodeDouble);
+    }
+  }
+
+  /**
+   * The full IEEE-754 ladder, in {@link Double#compare} order: signed zeros are distinguished, the
+   * infinities bound the finite range, and NaN sorts above everything.
+   */
+  @Test
+  public void testDoubleOrderingAcrossSpecialValues() {
+    List<Double> ascending =
+        Lists.newArrayList(
+            Double.NEGATIVE_INFINITY,
+            -Double.MAX_VALUE,
+            -1.0d,
+            -Double.MIN_VALUE,
+            -0.0d,
+            0.0d,
+            Double.MIN_VALUE,
+            1.0d,
+            Double.MAX_VALUE,
+            Double.POSITIVE_INFINITY,
+            Double.NaN);
+
+    assertOrderPreserved(ascending, TestZOrderByteUtil::encodeDouble);
+  }
+
+  @Test
+  public void testFloatOrderingAcrossSpecialValues() {
+    List<Float> ascending =
+        Lists.newArrayList(
+            Float.NEGATIVE_INFINITY,
+            -Float.MAX_VALUE,
+            -1.0f,
+            -Float.MIN_VALUE,
+            -0.0f,
+            0.0f,
+            Float.MIN_VALUE,
+            1.0f,
+            Float.MAX_VALUE,
+            Float.POSITIVE_INFINITY,
+            Float.NaN);
+
+    assertOrderPreserved(ascending, TestZOrderByteUtil::encodeFloat);
+  }
+
+  /**
+   * The encoding goes through {@link Double#doubleToLongBits}, which collapses every NaN to the
+   * canonical quiet NaN. Distinct NaN payloads must therefore produce identical bytes, otherwise
+   * the z-order key would not be deterministic for NaN.
+   */
+  @Test
+  public void testDoubleOrderedBytesCanonicalizesNaN() {
+    byte[] canonical = encodeDouble(Double.NaN);
+
+    for (long rawNaNBits :
+        new long[] {
+          0x7ff8000000000001L, // quiet NaN, non-zero payload
+          0x7fffffffffffffffL, // quiet NaN, all payload bits set
+          0xfff8000000000000L, // NaN with the sign bit set
+          0x7ff0000000000001L // signalling NaN
+        }) {
+      double nan = Double.longBitsToDouble(rawNaNBits);
+      assertThat(Double.isNaN(nan)).isTrue();
+
+      byte[] actual = encodeDouble(nan);
+      assertThat(actual)
+          .as("NaN with raw bits 0x%016x must encode as the canonical NaN", rawNaNBits)
+          .isEqualTo(canonical);
+    }
+  }
+
   @Test
   public void testByteTruncatedOrFillNullIsZeroArray() {
     ByteBuffer buffer = ByteBuffer.allocate(128);
@@ -388,4 +515,40 @@ public void testByteTruncatedOrFillNullIsZeroArray() {
 
     assertThat(actualBytes).isEqualTo(expected.array());
   }
+
+  private static byte[] encodeDouble(double value) {
+    return ZOrderByteUtils.doubleToOrderedBytes(value, ZOrderByteUtils.allocatePrimitiveBuffer())
+        .array();
+  }
+
+  private static byte[] encodeFloat(float value) {
+    return ZOrderByteUtils.floatToOrderedBytes(value, ZOrderByteUtils.allocatePrimitiveBuffer())
+        .array();
+  }
+
+  /**
+   * Asserts that the given values, which must already be in ascending order, encode to
+   * lexicographically ascending bytes.
+   */
+  private static <T extends Comparable<T>> void assertOrderPreserved(
+      List<T> ascendingValues, Function<T, byte[]> encoder) {
+    for (int i = 1; i < ascendingValues.size(); i++) {
+      T smaller = ascendingValues.get(i - 1);
+      T larger = ascendingValues.get(i);
+      assertThat(smaller).isLessThan(larger);
+
+      byte[] smallerBytes = encoder.apply(smaller);
+      byte[] largerBytes = encoder.apply(larger);
+
+      assertThat(UnsignedBytes.lexicographicalComparator().compare(smallerBytes, largerBytes))
+          .as(
+              "Ordering of %s should match ordering of bytes, %s -> %s is not less than %s -> %s",
+              smaller.getClass().getSimpleName(),
+              smaller,
+              Arrays.toString(smallerBytes),
+              larger,
+              Arrays.toString(largerBytes))
+          .isNegative();
+    }
+  }
 }