[GH-2028] Implementations of GeographyCollection & ShapeIndexGeography on S2Geography Apache Sedona

common/src/main/java/org/apache/sedona/common/S2Geography/EncodeOptions.java

@@ -36,6 +36,11 @@ public enum CodingHint {
 
   public EncodeOptions() {}
 
+  public EncodeOptions(EncodeOptions other) {
+    this.codingHint = other.codingHint;
+    this.enableLazyDecode = other.enableLazyDecode;
+    this.includeCovering = other.includeCovering;
+  }
   /** Control FAST vs. COMPACT encoding. */
   public void setCodingHint(CodingHint hint) {
     this.codingHint = hint;

common/src/main/java/org/apache/sedona/common/S2Geography/EncodedShapeIndexGeography.java

@@ -0,0 +1,140 @@
+/*
+ * 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.sedona.common.S2Geography;
+
+import com.esotericsoftware.kryo.io.Input;
+import com.esotericsoftware.kryo.io.UnsafeInput;
+import com.esotericsoftware.kryo.io.UnsafeOutput;
+import com.google.common.geometry.*;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.logging.Logger;
+
+public class EncodedShapeIndexGeography extends S2Geography {
+  private static final Logger logger = Logger.getLogger(EncodedShapeIndexGeography.class.getName());
+
+  public S2ShapeIndex shapeIndex;
+
+  /** Build an empty ShapeIndexGeography. */
+  public EncodedShapeIndexGeography() {
+    super(GeographyKind.ENCODED_SHAPE_INDEX);
+    this.shapeIndex = new S2ShapeIndex();
+  }
+
+  @Override
+  public int dimension() {
+    return -1;
+  }
+
+  @Override
+  public int numShapes() {
+    return shapeIndex.getShapes().size();
+  }
+
+  @Override
+  public S2Shape shape(int id) {
+    return shapeIndex.getShapes().get(id);
+  }
+
+  @Override
+  public S2Region region() {
+    return new S2ShapeIndexRegion(shapeIndex);
+  }
+
+  /**
+   * Index every S2Shape from the given Geography.
+   *
+   * @return the last shapeId assigned.
+   */
+  public int addIndex(S2Geography geog) {
+    int lastId = -1;
+    for (int i = 0, n = geog.numShapes(); i < n; i++) {
+      shapeIndex.add(geog.shape(i));
+      // since add() appends to the end, its index is size-1
+      lastId = shapeIndex.getShapes().size() - 1;
+    }
+    return lastId;
+  }
+
+  /** Add one raw shape into the index, return its new ID */
+  public int addIndex(S2Shape shape) {
+    shapeIndex.add(shape);
+    return shapeIndex.getShapes().size() - 1;
+  }
+
+  @Override
+  protected void encode(UnsafeOutput os, EncodeOptions opts) throws IOException {
+    throw new IOException("Encode() not implemented for EncodedShapeIndexGeography()");
+  }
+  // decode
+  /** This is what decodeTagged() actually calls */
+  public static EncodedShapeIndexGeography decode(Input in, EncodeTag tag) throws IOException {
+    // cast to UnsafeInput—will work if you always pass a Kryo-backed stream
+    if (!(in instanceof UnsafeInput)) {
+      throw new IllegalArgumentException("Expected UnsafeInput");
+    }
+    return decode((UnsafeInput) in, tag);
+  }
+
+  public static EncodedShapeIndexGeography decode(UnsafeInput in, EncodeTag tag)
+      throws IOException {
+    EncodedShapeIndexGeography encodedShapeIndexGeography = new EncodedShapeIndexGeography();
+
+    // EMPTY
+    if ((tag.getFlags() & EncodeTag.FLAG_EMPTY) != 0) {
+      logger.fine("Decoded empty EncodedShapeIndexGeography.");
+      return encodedShapeIndexGeography;
+    }
+
+    // 2) Skip past any covering cell-IDs written by encodeTagged
+    tag.skipCovering(in);
+
+    int length = in.readInt();
+    if (length < 0) {
+      throw new IOException("Invalid payload length: " + length);
+    }
+
+    // 2) read exactly that many bytes
+    byte[] payload = new byte[length];
+    in.readBytes(payload, 0, length);
+    // 3) hand *only* those bytes to S2‐Coder via Bytes adapter
+    PrimitiveArrays.Bytes bytes =
+        new PrimitiveArrays.Bytes() {
+          @Override
+          public long length() {
+            return payload.length;
+          }
+
+          @Override
+          public byte get(long i) {
+            return payload[(int) i];
+          }
+        };
+
+    List<S2Shape> s2Shapes = new ArrayList<>();
+    if (tag.isCompact()) {
+      s2Shapes = VectorCoder.COMPACT_SHAPE.decode(bytes);
+    } else {
+      s2Shapes = VectorCoder.FAST_SHAPE.decode(bytes);
+    }
+    for (S2Shape shape : s2Shapes) encodedShapeIndexGeography.addIndex(shape);
+    return encodedShapeIndexGeography;
+  }
+}

common/src/main/java/org/apache/sedona/common/S2Geography/GeographyCollection.java

@@ -0,0 +1,147 @@
+/*
+ * 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.sedona.common.S2Geography;
+
+import com.esotericsoftware.kryo.io.UnsafeInput;
+import com.esotericsoftware.kryo.io.UnsafeOutput;
+import com.google.common.collect.ImmutableList;
+import com.google.common.geometry.*;
+import java.io.*;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.logging.Logger;
+
+/** A Geography wrapping zero or more Geography objects, representing a GEOMETRYCOLLECTION. */
+public class GeographyCollection extends S2Geography {
+  private static final Logger logger = Logger.getLogger(GeographyCollection.class.getName());
+
+  public final List<S2Geography> features;
+  public final List<Integer> numShapesList;
+  public int totalShapes;
+
+  /** Constructs an empty GeographyCollection. */
+  public GeographyCollection() {
+    super(GeographyKind.GEOGRAPHY_COLLECTION);
+    this.features = new ArrayList<>();
+    this.numShapesList = new ArrayList<>();
+    this.totalShapes = 0;
+  }
+
+  /** Wraps existing Geography features. */
+  public GeographyCollection(List<S2Geography> features) {
+    super(GeographyKind.GEOGRAPHY_COLLECTION);
+    this.features = new ArrayList<>(features);
+    this.numShapesList = new ArrayList<>();
+    this.totalShapes = 0;
+    countShapes();
+  }
+
+  @Override
+  public int dimension() {
+    // Mixed or empty → return -1; uniform → return 0,1,2
+    return computeDimensionFromShapes();
+  }
+
+  @Override
+  public int numShapes() {
+    return totalShapes;
+  }
+
+  @Override
+  public S2Shape shape(int id) {
+    int sum = 0;
+    for (int i = 0; i < features.size(); i++) {
+      int n = numShapesList.get(i);
+      sum += n;
+      if (id < sum) {
+        // index within this feature
+        return features.get(i).shape(id - (sum - n));
+      }
+    }
+    throw new IllegalArgumentException("Shape id out of bounds: " + id);
+  }
+
+  @Override
+  public S2Region region() {
+    Collection<S2Region> regs = new ArrayList<>();
+    for (S2Geography geo : features) {
+      regs.add(geo.region());
+    }
+    return new S2RegionUnion(regs);
+  }
+
+  /** Returns an immutable copy of the features list. */
+  public List<S2Geography> getFeatures() {
+    return ImmutableList.copyOf(features);
+  }
+
+  @Override
+  public void encode(UnsafeOutput out, EncodeOptions opts) throws IOException {
+    // Top-level collection encodes its size and then each child with tagging
+    // Never include coverings for children (only a top-level concept
+    EncodeOptions childOptions = new EncodeOptions(opts);
+    childOptions.setIncludeCovering(false);
+    out.writeInt(features.size());
+    for (S2Geography feature : features) {
+      feature.encodeTagged(out, opts);
+    }
+    out.flush();
+  }
+
+  /** Decodes a GeographyCollection from a tagged input stream. */
+  public static GeographyCollection decode(UnsafeInput in, EncodeTag tag) throws IOException {
+    GeographyCollection geo = new GeographyCollection();
+
+    // Handle EMPTY flag
+    if ((tag.getFlags() & EncodeTag.FLAG_EMPTY) != 0) {
+      logger.fine("Decoded empty GeographyCollection.");
+      return geo;
+    }
+
+    // Skip any covering data
+    tag.skipCovering(in);
+
+    // 3) Ensure we have at least 4 bytes for the count
+    if (in.available() < Integer.BYTES) {
+      throw new IOException("GeographyCollection.decodeTagged error: insufficient header bytes");
+    }
+
+    // Read feature count
+    int n = in.readInt();
+    for (int i = 0; i < n; i++) {
+      tag = EncodeTag.decode(in);
+      // avoid creating new stream, directly call S2Geography.decode
+      S2Geography feature = S2Geography.decode(in, tag);
+      geo.features.add(feature);
+    }
+    geo.countShapes();
+    return geo;
+  }
+
+  private void countShapes() {
+    numShapesList.clear();
+    totalShapes = 0;
+    for (S2Geography geo : features) {
+      int n = geo.numShapes();
+      numShapesList.add(n);
+      totalShapes += n;
+    }
+  }
+}