Requires pyarrow 21 for geometry type

Author: Gawaboumga

pyiceberg/avro/reader.py

@@ -260,6 +260,16 @@ def skip(self, decoder: BinaryDecoder) -> None:
         decoder.skip_bytes()
 
 
+class GeographyReader(BinaryReader):
+    """Reads a geography from the stream.
+    """
+
+
+class GeometryReader(BinaryReader):
+    """Reads a geometry from the stream.
+    """
+
+
 @dataclass(frozen=True, init=False)
 class DecimalReader(Reader):
     """Reads a value as a decimal.

pyiceberg/avro/resolver.py

@@ -35,6 +35,8 @@
     DoubleReader,
     FixedReader,
     FloatReader,
+    GeographyReader,
+    GeometryReader,
     IntegerReader,
     ListReader,
     MapReader,
@@ -60,6 +62,8 @@
     DoubleWriter,
     FixedWriter,
     FloatWriter,
+    GeographyWriter,
+    GeometryWriter,
     IntegerWriter,
     ListWriter,
     MapWriter,
@@ -96,6 +100,8 @@
     FloatType,
     IcebergType,
     IntegerType,
+    GeographyType,
+    GeometryType,
     ListType,
     LongType,
     MapType,
@@ -181,6 +187,12 @@ def visit_float(self, float_type: FloatType) -> Writer:
     def visit_double(self, double_type: DoubleType) -> Writer:
         return DoubleWriter()
 
+    def visit_geography(self, geography_type: GeographyType) -> Writer:
+        return GeographyWriter()
+
+    def visit_geometry(self, geometry_type: GeometryType) -> Writer:
+        return GeometryWriter()
+
     def visit_date(self, date_type: DateType) -> Writer:
         return DateWriter()
 
@@ -338,6 +350,12 @@ def visit_decimal(self, decimal_type: DecimalType, partner: Optional[IcebergType
     def visit_date(self, date_type: DateType, partner: Optional[IcebergType]) -> Writer:
         return DateWriter()
 
+    def visit_geography(self, geography_type: GeographyType, partner: Optional[IcebergType]) -> Writer:
+        return GeographyWriter()
+
+    def visit_geometry(self, geometry_type: GeometryType, partner: Optional[IcebergType]) -> Writer:
+        return GeometryWriter()
+
     def visit_time(self, time_type: TimeType, partner: Optional[IcebergType]) -> Writer:
         return TimeWriter()
 
@@ -504,6 +522,12 @@ def visit_fixed(self, fixed_type: FixedType, partner: Optional[IcebergType]) ->
     def visit_binary(self, binary_type: BinaryType, partner: Optional[IcebergType]) -> Reader:
         return BinaryReader()
 
+    def visit_geography(self, geography_type: GeographyType, partner: Optional[IcebergType]) -> Reader:
+        return GeographyReader()
+
+    def visit_geometry(self, geometry_type: GeometryType, partner: Optional[IcebergType]) -> Reader:
+        return GeometryReader()
+
     def visit_unknown(self, unknown_type: UnknownType, partner: Optional[IcebergType]) -> Reader:
         return UnknownReader()
 

pyiceberg/avro/writer.py

@@ -161,6 +161,18 @@ def write(self, encoder: BinaryEncoder, val: Any) -> None:
         encoder.write_bytes(val)
 
 
+@dataclass(frozen=True)
+class GeographyWriter(Writer):
+    def write(self, encoder: BinaryEncoder, val: Any) -> None:
+        encoder.write_bytes(val)
+
+
+@dataclass(frozen=True)
+class GeometryWriter(Writer):
+    def write(self, encoder: BinaryEncoder, val: Any) -> None:
+        encoder.write_bytes(val)
+
+
 @dataclass(frozen=True)
 class DecimalWriter(Writer):
     precision: int = dataclassfield()

pyiceberg/conversions.py

@@ -51,6 +51,8 @@
     DoubleType,
     FixedType,
     FloatType,
+    GeometryType,
+    GeographyType,
     IntegerType,
     LongType,
     PrimitiveType,
@@ -167,6 +169,8 @@ def _(_: UUIDType, value_str: str) -> uuid.UUID:
 
 @partition_to_py.register(FixedType)
 @partition_to_py.register(BinaryType)
+@partition_to_py.register(GeographyType)
+@partition_to_py.register(GeometryType)
 @handle_none
 def _(_: PrimitiveType, value_str: str) -> bytes:
     return bytes(value_str, UTF8)
@@ -275,6 +279,8 @@ def _(_: UUIDType, value: Union[uuid.UUID, bytes]) -> bytes:
 
 @to_bytes.register(BinaryType)
 @to_bytes.register(FixedType)
+@to_bytes.register(GeographyType)
+@to_bytes.register(GeometryType)
 def _(_: PrimitiveType, value: bytes) -> bytes:
     return value
 
@@ -355,6 +361,8 @@ def _(_: StringType, b: bytes) -> str:
 
 @from_bytes.register(BinaryType)
 @from_bytes.register(FixedType)
+@from_bytes.register(GeographyType)
+@from_bytes.register(GeometryType)
 @from_bytes.register(UUIDType)
 def _(_: PrimitiveType, b: bytes) -> bytes:
     return b
@@ -453,6 +461,8 @@ def _(t: FixedType, b: bytes) -> str:
 
 
 @to_json.register(BinaryType)
+@to_json.register(GeographyType)
+@to_json.register(GeometryType)
 def _(_: BinaryType, b: bytes) -> str:
     """Python bytes serializes into hexadecimal encoded string."""
     return codecs.encode(b, "hex").decode(UTF8)
@@ -580,6 +590,16 @@ def _(_: BinaryType, val: Union[bytes, str]) -> bytes:
         return val
 
 
+@from_json.register(GeographyType)
+def _(_: GeographyType, val: bytes) -> bytes:
+    return val
+
+
+@from_json.register(GeometryType)
+def _(_: GeometryType, val: bytes) -> bytes:
+    return val
+
+
 @from_json.register(DecimalType)
 def _(_: DecimalType, val: str) -> Decimal:
     """Convert JSON string into a Python Decimal."""

pyiceberg/io/pyarrow.py

@@ -164,6 +164,8 @@
     DoubleType,
     FixedType,
     FloatType,
+    GeographyType,
+    GeometryType,
     IcebergType,
     IntegerType,
     ListType,
@@ -802,6 +804,12 @@ def visit_unknown(self, _: UnknownType) -> pa.DataType:
     def visit_binary(self, _: BinaryType) -> pa.DataType:
         return pa.large_binary()
 
+    def visit_geography(self, _: GeographyType) -> pa.DataType:
+        return pa.large_binary()
+
+    def visit_geometry(self, _: GeometryType) -> pa.DataType:
+        return pa.large_binary()
+
 
 def _convert_scalar(value: Any, iceberg_type: IcebergType) -> pa.scalar:
     if not isinstance(iceberg_type, PrimitiveType):
@@ -2064,6 +2072,12 @@ def visit_uuid(self, uuid_type: UUIDType) -> str:
     def visit_binary(self, binary_type: BinaryType) -> str:
         return "BYTE_ARRAY"
 
+    def visit_geography(self, geography_type: GeographyType) -> str:
+        return "BYTE_ARRAY"
+
+    def visit_geometry(self, geometry_type: GeometryType) -> str:
+        return "BYTE_ARRAY"
+
     def visit_unknown(self, unknown_type: UnknownType) -> str:
         return "UNKNOWN"
 

pyiceberg/schema.py

@@ -48,6 +48,8 @@
     DoubleType,
     FixedType,
     FloatType,
+    GeographyType,
+    GeometryType,
     IcebergType,
     IntegerType,
     ListType,
@@ -556,6 +558,10 @@ def primitive(self, primitive: PrimitiveType, primitive_partner: Optional[P]) ->
             return self.visit_fixed(primitive, primitive_partner)
         elif isinstance(primitive, BinaryType):
             return self.visit_binary(primitive, primitive_partner)
+        elif isinstance(primitive, GeographyType):
+            return self.visit_geography(primitive, primitive_partner)
+        elif isinstance(primitive, GeometryType):
+            return self.visit_geometry(primitive, primitive_partner)
         elif isinstance(primitive, UnknownType):
             return self.visit_unknown(primitive, primitive_partner)
         else:
@@ -625,6 +631,14 @@ def visit_fixed(self, fixed_type: FixedType, partner: Optional[P]) -> T:
     def visit_binary(self, binary_type: BinaryType, partner: Optional[P]) -> T:
         """Visit a BinaryType."""
 
+    @abstractmethod
+    def visit_geography(self, date_type: GeographyType, partner: Optional[P]) -> T:
+        """Visit a GeographyType."""
+
+    @abstractmethod
+    def visit_geometry(self, date_type: GeometryType, partner: Optional[P]) -> T:
+        """Visit a GeometryType."""
+
     @abstractmethod
     def visit_unknown(self, unknown_type: UnknownType, partner: Optional[P]) -> T:
         """Visit a UnknownType."""
@@ -750,6 +764,10 @@ def primitive(self, primitive: PrimitiveType) -> T:
             return self.visit_uuid(primitive)
         elif isinstance(primitive, BinaryType):
             return self.visit_binary(primitive)
+        elif isinstance(primitive, GeographyType):
+            return self.visit_geography(primitive)
+        elif isinstance(primitive, GeometryType):
+            return self.visit_geometry(primitive)
         elif isinstance(primitive, UnknownType):
             return self.visit_unknown(primitive)
         else:
@@ -819,6 +837,14 @@ def visit_uuid(self, uuid_type: UUIDType) -> T:
     def visit_binary(self, binary_type: BinaryType) -> T:
         """Visit a BinaryType."""
 
+    @abstractmethod
+    def visit_geography(self, geography_type: GeographyType) -> T:
+        """Visit a GeographyType."""
+
+    @abstractmethod
+    def visit_geometry(self, geometry_type: GeometryType) -> T:
+        """Visit a GeometryType."""
+
     @abstractmethod
     def visit_unknown(self, unknown_type: UnknownType) -> T:
         """Visit a UnknownType."""
@@ -1665,9 +1691,13 @@ def _(file_type: StringType, read_type: IcebergType) -> IcebergType:
 
 
 @promote.register(BinaryType)
+@promote.register(GeographyType)
+@promote.register(GeometryType)
 def _(file_type: BinaryType, read_type: IcebergType) -> IcebergType:
     if isinstance(read_type, StringType):
         return read_type
+    elif isinstance(read_type, (GeographyType, GeometryType)):
+        return read_type
     else:
         raise ResolveError(f"Cannot promote an binary to {read_type}")
 

pyiceberg/transforms.py

@@ -72,6 +72,8 @@
     DateType,
     DecimalType,
     FixedType,
+    GeographyType,
+    GeometryType,
     IcebergType,
     IntegerType,
     LongType,
@@ -368,7 +370,7 @@ def hash_func(v: Any) -> int:
             def hash_func(v: Any) -> int:
                 return mmh3.hash(decimal_to_bytes(v))
 
-        elif isinstance(source, (StringType, FixedType, BinaryType)):
+        elif isinstance(source, (StringType, FixedType, BinaryType, GeographyType, GeometryType)):
 
             def hash_func(v: Any) -> int:
                 return mmh3.hash(v)