GeoDataFrame init

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/Integration.kt

@@ -57,7 +57,6 @@ import org.jetbrains.kotlinx.jupyter.api.libraries.JupyterIntegration
 import org.jetbrains.kotlinx.jupyter.api.libraries.resources
 import kotlin.reflect.KClass
 import kotlin.reflect.KProperty
-import kotlin.reflect.KType
 import kotlin.reflect.full.isSubtypeOf
 
 /** Users will get an error if their Kotlin Jupyter kernel is older than this version. */
@@ -70,29 +69,6 @@ internal class Integration(private val notebook: Notebook, private val options:
 
     val version = options["v"]
 
-    private fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
-        val code = codeWithConverter.with(argument)
-        return if (code.isNotBlank()) {
-            val result = execute(code)
-            if (codeWithConverter.hasConverter) {
-                result.name
-            } else {
-                null
-            }
-        } else {
-            null
-        }
-    }
-
-    private fun KotlinKernelHost.execute(
-        codeWithConverter: CodeWithConverter,
-        property: KProperty<*>,
-        type: KType,
-    ): VariableName? {
-        val variableName = "(${property.name}${if (property.returnType.isMarkedNullable) "!!" else ""} as $type)"
-        return execute(codeWithConverter, variableName)
-    }
-
     private fun KotlinKernelHost.updateImportDataSchemaVariable(
         importDataSchema: ImportDataSchema,
         property: KProperty<*>,

core/src/main/kotlin/org/jetbrains/kotlinx/dataframe/jupyter/execute.kt

@@ -0,0 +1,30 @@
+package org.jetbrains.kotlinx.dataframe.jupyter
+
+import org.jetbrains.kotlinx.dataframe.codeGen.CodeWithConverter
+import org.jetbrains.kotlinx.jupyter.api.KotlinKernelHost
+import org.jetbrains.kotlinx.jupyter.api.VariableName
+import kotlin.reflect.KProperty
+import kotlin.reflect.KType
+
+internal fun KotlinKernelHost.execute(codeWithConverter: CodeWithConverter, argument: String): VariableName? {
+    val code = codeWithConverter.with(argument)
+    return if (code.isNotBlank()) {
+        val result = execute(code)
+        if (codeWithConverter.hasConverter) {
+            result.name
+        } else {
+            null
+        }
+    } else {
+        null
+    }
+}
+
+internal fun KotlinKernelHost.execute(
+    codeWithConverter: CodeWithConverter,
+    property: KProperty<*>,
+    type: KType,
+): VariableName? {
+    val variableName = "(${property.name}${if (property.returnType.isMarkedNullable) "!!" else ""} as $type)"
+    return execute(codeWithConverter, variableName)
+}

dataframe-geo/build.gradle.kts

@@ -0,0 +1,69 @@
+import org.jetbrains.kotlin.gradle.tasks.BaseKotlinCompile
+import org.jetbrains.kotlin.gradle.tasks.KotlinCompile
+
+plugins {
+    with(libs.plugins) {
+        alias(kotlin.jvm)
+        alias(publisher)
+        alias(jupyter.api)
+        alias(ktlint)
+    }
+}
+
+group = "org.jetbrains.kotlinx"
+
+repositories {
+    // geo repository should come before Maven Central
+    maven("https://repo.osgeo.org/repository/release")
+    mavenCentral()
+}
+
+// https://stackoverflow.com/questions/26993105/i-get-an-error-downloading-javax-media-jai-core1-1-3-from-maven-central
+// jai core dependency should be excluded from geotools dependencies and added separately
+fun ExternalModuleDependency.excludeJaiCore() = exclude("javax.media", "jai_core")
+
+dependencies {
+    api(project(":core"))
+
+    implementation(libs.geotools.main) { excludeJaiCore() }
+    implementation(libs.geotools.shapefile) { excludeJaiCore() }
+    implementation(libs.geotools.geojson) { excludeJaiCore() }
+    implementation(libs.geotools.referencing) { excludeJaiCore() }
+    implementation(libs.geotools.epsg.hsql) { excludeJaiCore() }
+
+    implementation(libs.jai.core)
+
+    implementation(libs.jts.core)
+    implementation(libs.jts.io.common)
+
+    implementation(libs.ktor.client.core)
+    implementation(libs.ktor.client.cio)
+    implementation(libs.ktor.client.content.negotiation)
+    implementation(libs.ktor.serialization.kotlinx.json)
+
+    testImplementation(kotlin("test"))
+}
+
+tasks.withType<KotlinCompile>().configureEach {
+    val friendModule = project(":core")
+    val jarTask = friendModule.tasks.getByName("jar") as Jar
+    val jarPath = jarTask.archiveFile.get().asFile.absolutePath
+    (this as BaseKotlinCompile).friendPaths.from(jarPath)
+}
+
+kotlinPublications {
+    publication {
+        publicationName = "dataframeGeo"
+        artifactId = "dataframe-geo"
+        description = "GeoDataFrame API"
+        packageName = artifactId
+    }
+}
+
+tasks.processJupyterApiResources {
+    libraryProducers = listOf("org.jetbrains.kotlinx.dataframe.jupyter.IntegrationGeo")
+}
+
+tasks.test {
+    useJUnitPlatform()
+}

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

@@ -0,0 +1,69 @@
+package org.jetbrains.kotlinx.dataframe.geo
+
+import org.geotools.api.referencing.crs.CoordinateReferenceSystem
+import org.geotools.geometry.jts.JTS
+import org.geotools.referencing.CRS
+import org.jetbrains.kotlinx.dataframe.DataFrame
+import org.jetbrains.kotlinx.dataframe.api.update
+import org.jetbrains.kotlinx.dataframe.api.with
+
+/**
+ * A data structure representing a geographical DataFrame, combining spatial data with
+ * an optional Coordinate Reference System (CRS).
+ *
+ * @param T The type parameter extending `WithGeometry`, indicating the presence of a geometry column.
+ * @property df The underlying `DataFrame` containing geometries.
+ * @property crs The coordinate reference system associated with the data, if any.
+ */
+class GeoDataFrame<T : WithGeometry>(val df: DataFrame<T>, val crs: CoordinateReferenceSystem?) {
+    /**
+     * Creates a new `GeoDataFrame` using a specified transformation block on the 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.
+     */
+    fun update(updateBlock: DataFrame<T>.() -> DataFrame<T>): GeoDataFrame<T> = GeoDataFrame(df.updateBlock(), crs)
+
+    /**
+     * Transforms the geometries to a specified Coordinate Reference System (CRS).
+     *
+     * This function reprojects the geometry data from the current CRS to a target CRS.
+     * If no target CRS is specified and the `GeoDataFrame` has no CRS, WGS 84 is used by default.
+     *
+     * @param targetCrs The target CRS for transformation.
+     * @return A new `GeoDataFrame` with reprojected geometries and the specified CRS.
+     */
+    fun applyCrs(targetCrs: CoordinateReferenceSystem): GeoDataFrame<T> {
+        if (crs == null) {
+            return GeoDataFrame(df, targetCrs)
+        }
+        if (targetCrs == this.crs) return this
+        // Use WGS 84 by default TODO
+        val sourceCRS: CoordinateReferenceSystem = this.crs
+        val transform = CRS.findMathTransform(sourceCRS, targetCrs, true)
+        return GeoDataFrame(
+            df.update { geometry }.with { JTS.transform(it, transform) },
+            targetCrs,
+        )
+    }
+
+    override fun equals(other: Any?): Boolean {
+        if (this === other) return true
+        if (other !is GeoDataFrame<*>) 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)
+            }
+    }
+
+    override fun toString(): String = super.toString()
+
+    override fun hashCode(): Int = super.hashCode()
+
+    companion object {
+        val DEFAULT_CRS = CRS.decode("EPSG:4326", true)
+    }
+}

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

@@ -0,0 +1,38 @@
+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.MultiPolygon
+import org.locationtech.jts.geom.Polygon
+
+@DataSchema
+interface WithGeometry {
+    val geometry: Geometry
+}
+
+@DataSchema
+interface WithPolygon : WithGeometry {
+    override val geometry: Polygon
+}
+
+@DataSchema
+interface WithMultiPolygon : WithGeometry {
+    override val geometry: MultiPolygon
+}
+
+@Suppress("UNCHECKED_CAST")
+@get:JvmName("geometry")
+val ColumnsContainer<WithGeometry>.geometry: DataColumn<Geometry>
+    get() = get("geometry") as DataColumn<Geometry>
+
+@Suppress("UNCHECKED_CAST")
+@get:JvmName("geometryPolygon")
+val ColumnsContainer<WithPolygon>.geometry: DataColumn<Polygon>
+    get() = get("geometry") as DataColumn<Polygon>
+
+@Suppress("UNCHECKED_CAST")
+@get:JvmName("geometryMultiPolygon")
+val ColumnsContainer<WithMultiPolygon>.geometry: DataColumn<MultiPolygon>
+    get() = get("geometry") as DataColumn<MultiPolygon>

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

@@ -0,0 +1,15 @@
+package org.jetbrains.kotlinx.dataframe.geo
+
+import org.geotools.geometry.jts.ReferencedEnvelope
+import org.jetbrains.kotlinx.dataframe.api.asIterable
+import org.jetbrains.kotlinx.dataframe.geo.jts.computeBounds
+
+/**
+ * Computes the bounding envelope for all geometries in a `GeoDataFrame`,
+ * considering the specified coordinate reference system (CRS).
+ *
+ * @receiver The `GeoDataFrame` containing the geometries for which to compute bounds.
+ * @return The bounding envelope that includes all geometries,
+ * associated with the CRS of the `GeoDataFrame`.
+ */
+fun GeoDataFrame<*>.bounds(): ReferencedEnvelope = ReferencedEnvelope(df.geometry.asIterable().computeBounds(), crs)

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

@@ -0,0 +1,126 @@
+package org.jetbrains.kotlinx.dataframe.geo.geocode
+
+import io.ktor.client.HttpClient
+import io.ktor.client.engine.cio.CIO
+import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
+import io.ktor.client.request.post
+import io.ktor.client.request.setBody
+import io.ktor.client.statement.bodyAsText
+import io.ktor.http.ContentType
+import io.ktor.http.contentType
+import io.ktor.serialization.kotlinx.json.json
+import kotlinx.coroutines.runBlocking
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.jsonArray
+import kotlinx.serialization.json.jsonObject
+import kotlinx.serialization.json.jsonPrimitive
+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
+
+object Geocoder {
+
+    private val url = "https://geo2.datalore.jetbrains.com/map_data/geocoding"
+
+    private fun countryQuery(country: String) =
+        """
+         {
+        "region_query_names" : [ "$country" ],
+        "region_query_countries" : null,
+        "region_query_states" : null,
+        "region_query_counties" : null,
+        "ambiguity_resolver" : {
+          "ambiguity_resolver_ignoring_strategy" : null,
+          "ambiguity_resolver_box" : null,
+          "ambiguity_resolver_closest_coord" : null
+        }
+        }
+        """.trimIndent()
+
+    private fun geocodeQuery(countries: List<String>) =
+        """
+{
+  "version" : 3,
+  "mode" : "by_geocoding",
+  "feature_options" : [ "limit", "position", "centroid" ],
+  "resolution" : null,
+  "view_box" : null,
+  "fetched_ids" : null,
+  "region_queries" : [ 
+  ${countries.joinToString(",\n") { countryQuery(it) }}
+   ],
+  "scope" : [ ],
+  "level" : "country",
+  "namesake_example_limit" : 10,
+  "allow_ambiguous" : false
+}
+        """.trimIndent()
+
+    private fun idsQuery(ids: List<String>) =
+        """
+        {"version": 3, 
+        "mode": "by_id", 
+        "feature_options": ["boundary"], 
+        "resolution": 5, 
+        "view_box": null, 
+        "fetched_ids": null, 
+        "ids": [${ids.joinToString(", ") { "\"" + it + "\"" }}]}
+        """.trimIndent()
+
+    private val client = HttpClient(CIO) {
+        install(ContentNegotiation) {
+            json(
+                Json {
+                    prettyPrint = true
+                    isLenient = true
+                },
+            )
+        }
+    }
+
+    fun geocodeCountries(countries: List<String>): GeoDataFrame<*> {
+        val query = geocodeQuery(countries)
+        val foundNames = mutableListOf<String>()
+        val geometries = mutableListOf<Geometry>()
+        runBlocking {
+            val responseString = client.post(url) {
+                contentType(ContentType.Application.Json)
+                // headers[HttpHeaders.AcceptEncoding] = "gzip"
+                setBody(query)
+            }.bodyAsText()
+            val ids = mutableListOf<String>()
+
+            Json.parseToJsonElement(responseString).jsonObject["data"]!!.jsonObject["answers"]!!.jsonArray.forEach {
+                it.jsonObject["features"]!!.jsonArray.single().jsonObject.also {
+                    foundNames.add(it["name"]!!.jsonPrimitive.content)
+                    ids.add(it["id"]!!.jsonPrimitive.content)
+                }
+            }
+            val idsQuery = idsQuery(ids)
+
+            val responseStringGeometries = client.post(url) {
+                contentType(ContentType.Application.Json)
+                // headers[HttpHeaders.AcceptEncoding] = "gzip"
+                setBody(idsQuery)
+            }.bodyAsText()
+
+            val geoJsonReader = GeoJsonReader(GeometryFactory())
+            Json.parseToJsonElement(
+                responseStringGeometries,
+            ).jsonObject["data"]!!.jsonObject["answers"]!!.jsonArray.forEach {
+                it.jsonObject["features"]!!.jsonArray.single().jsonObject.also {
+                    val boundary = it["boundary"]!!.jsonPrimitive.content
+                    geometries.add(geoJsonReader.read(boundary))
+                }
+            }
+        }
+        return dataFrameOf(
+            "country" to countries,
+            "foundName" to foundNames,
+            "geometry" to geometries,
+        ).toGeo()
+    }
+}