Parquet: Compute geometry bounding box metrics

api/src/main/java/org/apache/iceberg/geospatial/WKBBoundingBox.java

@@ -0,0 +1,217 @@
+/*
+ * 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.
+ */
+package org.apache.iceberg.geospatial;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+
+/**
+ * Computes a two-dimensional (XY) bounding box from geometry values encoded as Well-Known Binary
+ * (WKB).
+ *
+ * <p>The bounding box is the minimum axis-aligned rectangle that contains every coordinate of the
+ * geometries. Only the X and Y dimensions are considered; any Z or M values present in the WKB are
+ * read past and ignored. Coordinates whose X or Y is {@code NaN} are skipped, matching the spec
+ * rule that null or NaN values in a coordinate dimension do not contribute to bounds (an empty
+ * geometry such as {@code POINT EMPTY} therefore contributes nothing).
+ *
+ * <p>Parsing follows the OGC Simple Feature Access WKB layout and validates the input defensively:
+ * a malformed or truncated buffer results in an {@link IllegalArgumentException} rather than an
+ * out-of-bounds read.
+ */
+public class WKBBoundingBox {
+
+  // OGC WKB base geometry type codes (after stripping the dimension offset).
+  private static final int TYPE_POINT = 1;
+  private static final int TYPE_LINE_STRING = 2;
+  private static final int TYPE_POLYGON = 3;
+  private static final int TYPE_MULTI_POINT = 4;
+  private static final int TYPE_MULTI_LINE_STRING = 5;
+  private static final int TYPE_MULTI_POLYGON = 6;
+  private static final int TYPE_GEOMETRY_COLLECTION = 7;
+
+  // Bounds the recursion depth for nested collections to reject pathological or malicious input.
+  private static final int MAX_DEPTH = 100;
+
+  private WKBBoundingBox() {}
+
+  /**
+   * Parses one WKB geometry and folds all of its XY coordinates into the given accumulator.
+   *
+   * <p>The buffer is read via a duplicate, so the caller's position and limit are left unchanged.
+   *
+   * @param wkb a buffer containing a single WKB geometry
+   * @param accumulator the accumulator to update with the geometry's coordinates
+   * @throws IllegalArgumentException if the buffer is not a well-formed WKB geometry
+   */
+  public static void accumulate(ByteBuffer wkb, XYAccumulator accumulator) {
+    Preconditions.checkArgument(wkb != null, "Invalid WKB buffer: null");
+    Preconditions.checkArgument(accumulator != null, "Invalid accumulator: null");
+    ByteBuffer buffer = wkb.duplicate();
+    parseGeometry(buffer, accumulator, 0);
+  }
+
+  private static void parseGeometry(ByteBuffer buffer, XYAccumulator accumulator, int depth) {
+    Preconditions.checkArgument(depth <= MAX_DEPTH, "Invalid WKB: nesting too deep");
+    checkRemaining(buffer, 5);
+
+    byte order = buffer.get();
+    if (order == 0) {
+      buffer.order(ByteOrder.BIG_ENDIAN);
+    } else if (order == 1) {
+      buffer.order(ByteOrder.LITTLE_ENDIAN);
+    } else {
+      throw new IllegalArgumentException("Invalid WKB byte order: " + order);
+    }
+
+    long typeCode = buffer.getInt() & 0xFFFFFFFFL;
+    int geometryType = (int) (typeCode % 1000);
+    int dimensionGroup = (int) (typeCode / 1000);
+    int numDimensions = numDimensions(dimensionGroup, typeCode);
+
+    switch (geometryType) {
+      case TYPE_POINT:
+        readCoordinate(buffer, accumulator, numDimensions);
+        break;
+      case TYPE_LINE_STRING:
+        readCoordinateSequence(buffer, accumulator, numDimensions);
+        break;
+      case TYPE_POLYGON:
+        int numRings = readCount(buffer);
+        for (int i = 0; i < numRings; i += 1) {
+          readCoordinateSequence(buffer, accumulator, numDimensions);
+        }
+        break;
+      case TYPE_MULTI_POINT:
+      case TYPE_MULTI_LINE_STRING:
+      case TYPE_MULTI_POLYGON:
+      case TYPE_GEOMETRY_COLLECTION:
+        int numElements = readCount(buffer);
+        for (int i = 0; i < numElements; i += 1) {
+          // Each child carries its own byte-order and type header.
+          parseGeometry(buffer, accumulator, depth + 1);
+        }
+        break;
+      default:
+        throw new IllegalArgumentException("Invalid or unsupported WKB geometry type: " + typeCode);
+    }
+  }
+
+  private static int numDimensions(int dimensionGroup, long typeCode) {
+    switch (dimensionGroup) {
+      case 0: // XY
+        return 2;
+      case 1: // XYZ
+      case 2: // XYM
+        return 3;
+      case 3: // XYZM
+        return 4;
+      default:
+        throw new IllegalArgumentException("Invalid or unsupported WKB geometry type: " + typeCode);
+    }
+  }
+
+  private static void readCoordinateSequence(
+      ByteBuffer buffer, XYAccumulator accumulator, int numDimensions) {
+    int numPoints = readCount(buffer);
+    // Validate the full extent up front so a bogus point count cannot drive a long read loop, but
+    // never pre-allocate from the count itself.
+    checkRemaining(buffer, (long) numPoints * numDimensions * Double.BYTES);
+    for (int i = 0; i < numPoints; i += 1) {
+      readCoordinate(buffer, accumulator, numDimensions);
+    }
+  }
+
+  private static void readCoordinate(
+      ByteBuffer buffer, XYAccumulator accumulator, int numDimensions) {
+    checkRemaining(buffer, (long) numDimensions * Double.BYTES);
+    double xCoord = buffer.getDouble();
+    double yCoord = buffer.getDouble();
+    // Skip any Z and/or M values; only X and Y contribute to the box.
+    for (int i = 2; i < numDimensions; i += 1) {
+      buffer.getDouble();
+    }
+    accumulator.addXY(xCoord, yCoord);
+  }
+
+  private static int readCount(ByteBuffer buffer) {
+    checkRemaining(buffer, Integer.BYTES);
+    long count = buffer.getInt() & 0xFFFFFFFFL;
+    Preconditions.checkArgument(count <= Integer.MAX_VALUE, "Invalid WKB element count: %s", count);
+    return (int) count;
+  }
+
+  private static void checkRemaining(ByteBuffer buffer, long bytes) {
+    Preconditions.checkArgument(
+        buffer.remaining() >= bytes, "Invalid WKB: unexpected end of buffer");
+  }
+
+  /**
+   * A mutable accumulator of the minimum and maximum X and Y coordinates seen so far.
+   *
+   * <p>Coordinates whose X or Y is {@code NaN} are ignored. An accumulator that has seen no non-NaN
+   * coordinate reports {@link #hasBounds()} as {@code false} and returns {@code null} bounds.
+   */
+  public static class XYAccumulator {
+    private double minX = Double.POSITIVE_INFINITY;
+    private double minY = Double.POSITIVE_INFINITY;
+    private double maxX = Double.NEGATIVE_INFINITY;
+    private double maxY = Double.NEGATIVE_INFINITY;
+    private boolean hasBounds = false;
+
+    /**
+     * Folds a single coordinate into the box, ignoring it if either dimension is {@code NaN}.
+     *
+     * @param xCoord the X coordinate
+     * @param yCoord the Y coordinate
+     */
+    public void addXY(double xCoord, double yCoord) {
+      if (Double.isNaN(xCoord) || Double.isNaN(yCoord)) {
+        return;
+      }
+      minX = Math.min(minX, xCoord);
+      minY = Math.min(minY, yCoord);
+      maxX = Math.max(maxX, xCoord);
+      maxY = Math.max(maxY, yCoord);
+      hasBounds = true;
+    }
+
+    /** Returns whether any non-NaN coordinate has been accumulated. */
+    public boolean hasBounds() {
+      return hasBounds;
+    }
+
+    /**
+     * Returns the lower corner (minimum X and Y) of the box, or {@code null} if no coordinate has
+     * been accumulated.
+     */
+    public GeospatialBound minBound() {
+      return hasBounds ? GeospatialBound.createXY(minX, minY) : null;
+    }
+
+    /**
+     * Returns the upper corner (maximum X and Y) of the box, or {@code null} if no coordinate has
+     * been accumulated.
+     */
+    public GeospatialBound maxBound() {
+      return hasBounds ? GeospatialBound.createXY(maxX, maxY) : null;
+    }
+  }
+}