geometry & other fixes

dataframe-geo/build.gradle.kts

@@ -7,6 +7,8 @@ plugins {
         alias(publisher)
         alias(jupyter.api)
         alias(ktlint)
+        alias(dataframe)
+        alias(ksp)
     }
 }
 

dataframe-geo/src/main/kotlin/org/jetbrains/kotlinx/dataframe/geo/GeoDataFrame.kt

@@ -17,12 +17,12 @@ import org.jetbrains.kotlinx.dataframe.api.with
  */
 class GeoDataFrame<T : WithGeometry>(val df: DataFrame<T>, val crs: CoordinateReferenceSystem?) {
     /**
-     * Creates a new `GeoDataFrame` using a specified transformation block on the underlying DataFrame.
+     * Creates a new `GeoDataFrame` with the modified underlying DataFrame.
      *
-     * @param updateBlock The block defining the transformations to be applied to the DataFrame.
-     * @return A new `GeoDataFrame` instance with updated data and the same CRS.
+     * @param block The block defining the transformations to be applied to the DataFrame.
+     * @return A new `GeoDataFrame` instance with updated dataframe and the same CRS.
      */
-    fun update(updateBlock: DataFrame<T>.() -> DataFrame<T>): GeoDataFrame<T> = GeoDataFrame(df.updateBlock(), crs)
+    inline fun modify(block: DataFrame<T>.() -> DataFrame<T>): GeoDataFrame<T> = GeoDataFrame(df.block(), crs)
 
     /**
      * Transforms the geometries to a specified Coordinate Reference System (CRS).
@@ -49,19 +49,26 @@ class GeoDataFrame<T : WithGeometry>(val df: DataFrame<T>, val crs: CoordinateRe
 
     override fun equals(other: Any?): Boolean {
         if (this === other) return true
-        if (other !is GeoDataFrame<*>) return false
+        if (javaClass != other?.javaClass) return false
 
-        return df == other.df &&
-            when {
-                crs == null && other.crs == null -> true
-                crs == null || other.crs == null -> false
-                else -> CRS.equalsIgnoreMetadata(crs, other.crs)
-            }
+        other as GeoDataFrame<*>
+
+        if (df != other.df) return false
+
+        return when {
+            crs == null && other.crs == null -> true
+            crs == null || other.crs == null -> false
+            else -> CRS.equalsIgnoreMetadata(crs, other.crs)
+        }
     }
 
-    override fun toString(): String = super.toString()
+    override fun hashCode(): Int {
+        var result = df.hashCode()
+        result = 31 * result + (crs?.hashCode() ?: 0)
+        return result
+    }
 
-    override fun hashCode(): Int = super.hashCode()
+    override fun toString(): String = "GeoDataFrame(df=$df, crs=$crs)"
 
     companion object {
         val DEFAULT_CRS = CRS.decode("EPSG:4326", true)

dataframe-geo/src/main/kotlin/org/jetbrains/kotlinx/dataframe/geo/WithGeometry.kt

@@ -1,10 +1,12 @@
 package org.jetbrains.kotlinx.dataframe.geo
 
-import org.jetbrains.kotlinx.dataframe.ColumnsContainer
-import org.jetbrains.kotlinx.dataframe.DataColumn
 import org.jetbrains.kotlinx.dataframe.annotations.DataSchema
 import org.locationtech.jts.geom.Geometry
+import org.locationtech.jts.geom.LineString
+import org.locationtech.jts.geom.MultiLineString
+import org.locationtech.jts.geom.MultiPoint
 import org.locationtech.jts.geom.MultiPolygon
+import org.locationtech.jts.geom.Point
 import org.locationtech.jts.geom.Polygon
 
 @DataSchema
@@ -13,26 +15,31 @@ interface WithGeometry {
 }
 
 @DataSchema
-interface WithPolygon : WithGeometry {
+interface WithPolygonGeometry : WithGeometry {
     override val geometry: Polygon
 }
 
 @DataSchema
-interface WithMultiPolygon : WithGeometry {
+interface WithMultiPolygonGeometry : WithGeometry {
     override val geometry: MultiPolygon
 }
 
-@Suppress("UNCHECKED_CAST")
-@get:JvmName("geometry")
-val ColumnsContainer<WithGeometry>.geometry: DataColumn<Geometry>
-    get() = get("geometry") as DataColumn<Geometry>
+@DataSchema
+interface WithPointGeometry : WithGeometry {
+    override val geometry: Point
+}
 
-@Suppress("UNCHECKED_CAST")
-@get:JvmName("geometryPolygon")
-val ColumnsContainer<WithPolygon>.geometry: DataColumn<Polygon>
-    get() = get("geometry") as DataColumn<Polygon>
+@DataSchema
+interface WithMultiPointGeometry : WithGeometry {
+    override val geometry: MultiPoint
+}
 
-@Suppress("UNCHECKED_CAST")
-@get:JvmName("geometryMultiPolygon")
-val ColumnsContainer<WithMultiPolygon>.geometry: DataColumn<MultiPolygon>
-    get() = get("geometry") as DataColumn<MultiPolygon>
+@DataSchema
+interface WithLineStringGeometry : WithGeometry {
+    override val geometry: LineString
+}
+
+@DataSchema
+interface WithMultiLineStringGeometry : WithGeometry {
+    override val geometry: MultiLineString
+}

dataframe-geo/src/main/kotlin/org/jetbrains/kotlinx/dataframe/geo/geocode/Geocoder.kt

@@ -14,13 +14,18 @@ import kotlinx.serialization.json.Json
 import kotlinx.serialization.json.jsonArray
 import kotlinx.serialization.json.jsonObject
 import kotlinx.serialization.json.jsonPrimitive
+import org.jetbrains.annotations.ApiStatus.Experimental
 import org.jetbrains.kotlinx.dataframe.api.dataFrameOf
 import org.jetbrains.kotlinx.dataframe.geo.GeoDataFrame
 import org.jetbrains.kotlinx.dataframe.geo.toGeo
 import org.locationtech.jts.geom.Geometry
 import org.locationtech.jts.geom.GeometryFactory
 import org.locationtech.jts.io.geojson.GeoJsonReader
 
+/**
+ * Experimental geo coding utility.
+ */
+@Experimental
 object Geocoder {
 
     private val url = "https://geo2.datalore.jetbrains.com/map_data/geocoding"

dataframe-geo/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/IntegrationGeo.kt

@@ -4,8 +4,12 @@ package org.jetbrains.kotlinx.dataframe.jupyter
 
 import org.jetbrains.kotlinx.dataframe.geo.GeoDataFrame
 import org.jetbrains.kotlinx.dataframe.geo.WithGeometry
-import org.jetbrains.kotlinx.dataframe.geo.WithMultiPolygon
-import org.jetbrains.kotlinx.dataframe.geo.WithPolygon
+import org.jetbrains.kotlinx.dataframe.geo.WithLineStringGeometry
+import org.jetbrains.kotlinx.dataframe.geo.WithMultiLineStringGeometry
+import org.jetbrains.kotlinx.dataframe.geo.WithMultiPointGeometry
+import org.jetbrains.kotlinx.dataframe.geo.WithMultiPolygonGeometry
+import org.jetbrains.kotlinx.dataframe.geo.WithPointGeometry
+import org.jetbrains.kotlinx.dataframe.geo.WithPolygonGeometry
 import org.jetbrains.kotlinx.dataframe.impl.codeGen.ReplCodeGeneratorImpl
 import org.jetbrains.kotlinx.jupyter.api.FieldHandler
 import org.jetbrains.kotlinx.jupyter.api.FieldHandlerExecution
@@ -29,17 +33,27 @@ internal class IntegrationGeo : JupyterIntegration() {
         import("org.jetbrains.kotlinx.dataframe.geo.jts.*")
         import("org.jetbrains.kotlinx.dataframe.geo.geotools.*")
         import("org.jetbrains.kotlinx.dataframe.geo.geocode.*")
+        import("org.geotools.referencing.CRS")
+        import("org.locationtech.jts.geom.*")
         onLoaded {
             useSchema<WithGeometry>()
-            useSchema<WithPolygon>()
-            useSchema<WithMultiPolygon>()
+            useSchema<WithPolygonGeometry>()
+            useSchema<WithMultiPolygonGeometry>()
+            useSchema<WithPointGeometry>()
+            useSchema<WithMultiPointGeometry>()
+            useSchema<WithLineStringGeometry>()
+            useSchema<WithMultiLineStringGeometry>()
         }
         val replCodeGeneratorImpl = ReplCodeGeneratorImpl()
         replCodeGeneratorImpl.process(WithGeometry::class)
-        replCodeGeneratorImpl.process(WithPolygon::class)
-        replCodeGeneratorImpl.process(WithMultiPolygon::class)
+        replCodeGeneratorImpl.process(WithPolygonGeometry::class)
+        replCodeGeneratorImpl.process(WithMultiPolygonGeometry::class)
+        replCodeGeneratorImpl.process(WithPointGeometry::class)
+        replCodeGeneratorImpl.process(WithMultiPointGeometry::class)
+        replCodeGeneratorImpl.process(WithLineStringGeometry::class)
+        replCodeGeneratorImpl.process(WithMultiLineStringGeometry::class)
         val execution = FieldHandlerFactory.createUpdateExecution<GeoDataFrame<*>> { geo, kProperty ->
-            // TODO rewrite
+            // TODO rewrite better
             val generatedDf = execute(
                 codeWithConverter = replCodeGeneratorImpl.process(geo.df, kProperty),
                 "(${kProperty.name}.df as DataFrame<*>)",