Merge pull request #614 from databrickslabs/dbr17/ogr-readers

mjohns-databricks · web-flow · commit e79b0da10bcc · 2025-10-06T11:11:58.000-04:00
OGR Named Readers
diff --git a/scripts/docker/README-DOCKER.md b/scripts/docker/README-DOCKER.md
@@ -1,5 +1,5 @@
 For tests to run in docker we need to copy the jar to site-packages/pyspark/jars.
-Example: /root/mosaic/python/dev/lib/python3.12/site-packages/pyspark/jars/
+Example: /root/geobrix/python/dev/lib/python3.12/site-packages/pyspark/jars/
 spark.addArtifact does not work due to permissions issues.
 Also there could be annoying warnings about jypiter_client which isnt something relevant to the pytest we do.
 If those warnings occur use:  export JUPYTER_PLATFORM_DIRS=1
diff --git a/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister b/src/main/resources/META-INF/services/org.apache.spark.sql.sources.DataSourceRegister
@@ -1,4 +1,7 @@
 com.databricks.labs.gbx.rasterx.ds.gdal.GDAL_DataSource
 com.databricks.labs.gbx.vectorx.ds.ogr.OGR_DataSource
 com.databricks.labs.gbx.vectorx.ds.shp.ShapeFile_DataSource
+com.databricks.labs.gbx.vectorx.ds.gdb.FileGDB_DataSource
+com.databricks.labs.gbx.vectorx.ds.geojson.GeoJSON_DataSource
+com.databricks.labs.gbx.vectorx.ds.gpkg.GPKG_DataSource
 com.databricks.labs.gbx.ds.register.RegisterDataSource
diff --git a/src/main/scala/com/databricks/labs/gbx/ds/DataSourceExtras.scala b/src/main/scala/com/databricks/labs/gbx/ds/DataSourceExtras.scala
@@ -0,0 +1,23 @@
+package com.databricks.labs.gbx.ds
+
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+import scala.jdk.CollectionConverters._
+
+trait DataSourceExtras {
+
+    def dsExtraMap(checkMap: Map[String, String] = Map.empty): Map[String, String]
+
+    def extraJavaUtilMap(properties: java.util.Map[String, String]): java.util.Map[String, String] = {
+        val cMap = properties.asScala.toMap
+        val newMap = cMap ++ dsExtraMap(checkMap = cMap)
+        newMap.asJava
+    }
+
+    def extraCaseInsensitiveStringMap(options: CaseInsensitiveStringMap): CaseInsensitiveStringMap = {
+        val cMap = options.asCaseSensitiveMap().asScala.toMap
+        val newMap = cMap ++ dsExtraMap(checkMap = cMap)
+        new CaseInsensitiveStringMap(newMap.asJava)
+    }
+
+}
diff --git a/src/main/scala/com/databricks/labs/gbx/vectorx/ds/gdb/FileGDB_DataSource.scala b/src/main/scala/com/databricks/labs/gbx/vectorx/ds/gdb/FileGDB_DataSource.scala
@@ -0,0 +1,26 @@
+package com.databricks.labs.gbx.vectorx.ds.gdb
+
+import com.databricks.labs.gbx.ds.DataSourceExtras
+import com.databricks.labs.gbx.vectorx.ds.ogr.OGR_DataSource
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+//noinspection ScalaUnusedSymbol
+class FileGDB_DataSource extends OGR_DataSource with DataSourceExtras {
+
+    override def dsExtraMap(checkMap: Map[String, String] = Map.empty): Map[String, String] = Map(
+        "driverName" -> "OpenFileGDB"
+    )
+
+    override def shortName(): String = "file_gdb"
+
+    override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+        super.inferSchema(extraCaseInsensitiveStringMap(options))
+    }
+
+    override def getTable(schema: StructType, partitions: Array[Transform], properties: java.util.Map[String, String]): Table = {
+        super.getTable(schema, partitions, extraJavaUtilMap(properties))
+    }
+}
diff --git a/src/main/scala/com/databricks/labs/gbx/vectorx/ds/geojson/GeoJSON_DataSource.scala b/src/main/scala/com/databricks/labs/gbx/vectorx/ds/geojson/GeoJSON_DataSource.scala
@@ -0,0 +1,32 @@
+package com.databricks.labs.gbx.vectorx.ds.geojson
+
+import com.databricks.labs.gbx.ds.DataSourceExtras
+import com.databricks.labs.gbx.vectorx.ds.ogr.OGR_DataSource
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+
+//noinspection ScalaUnusedSymbol
+class GeoJSON_DataSource extends OGR_DataSource with DataSourceExtras{
+
+    // default to multi = true given common use
+    override def dsExtraMap(checkMap: Map[String, String] = Map.empty): Map[String, String] = {
+        if (checkMap.getOrElse("multi", "true").toBoolean) {
+            Map("driverName" -> "GeoJSONSeq")
+        } else {
+            Map("driverName" -> "GeoJSON")
+        }
+    }
+
+    override def shortName(): String = "geojson"
+
+    override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+        super.inferSchema(extraCaseInsensitiveStringMap(options))
+    }
+
+    override def getTable(schema: StructType, partitions: Array[Transform], properties: java.util.Map[String, String]): Table = {
+        super.getTable(schema, partitions, extraJavaUtilMap(properties))
+    }
+}
diff --git a/src/main/scala/com/databricks/labs/gbx/vectorx/ds/gpkg/GPKG_DataSource.scala b/src/main/scala/com/databricks/labs/gbx/vectorx/ds/gpkg/GPKG_DataSource.scala
@@ -0,0 +1,27 @@
+package com.databricks.labs.gbx.vectorx.ds.gpkg
+
+import com.databricks.labs.gbx.ds.DataSourceExtras
+import com.databricks.labs.gbx.vectorx.ds.ogr.OGR_DataSource
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
+
+//noinspection ScalaUnusedSymbol
+class GPKG_DataSource extends OGR_DataSource with DataSourceExtras{
+
+    override def dsExtraMap(checkMap: Map[String, String] = Map.empty): Map[String, String] = Map(
+        "driverName" -> "GPKG"
+    )
+
+    override def shortName(): String = "ogr_gpkg"
+
+    override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+        super.inferSchema(extraCaseInsensitiveStringMap(options))
+    }
+
+    override def getTable(schema: StructType, partitions: Array[Transform], properties: java.util.Map[String, String]): Table = {
+        super.getTable(schema, partitions, extraJavaUtilMap(properties))
+    }
+}
diff --git a/src/main/scala/com/databricks/labs/gbx/vectorx/ds/ogr/OGR_Batch.scala b/src/main/scala/com/databricks/labs/gbx/vectorx/ds/ogr/OGR_Batch.scala
@@ -15,7 +15,7 @@ class OGR_Batch(schema: StructType, options: Map[String, String]) extends Scan w
 
     override def planInputPartitions(): Array[InputPartition] = {
         val inPath = options("path")
-        val chunkSize = options("chunkSize").toInt
+        val chunkSize = options.getOrElse("chunkSize", "10000").toInt
         val driverName = options.getOrElse("driverName", "")
         val layerN = options.getOrElse("layerNumber", "0").toInt
         val layerName = options.getOrElse("layerName", "")
diff --git a/src/main/scala/com/databricks/labs/gbx/vectorx/ds/shp/ShapeFile_DataSource.scala b/src/main/scala/com/databricks/labs/gbx/vectorx/ds/shp/ShapeFile_DataSource.scala
@@ -1,10 +1,27 @@
 package com.databricks.labs.gbx.vectorx.ds.shp
 
+import com.databricks.labs.gbx.ds.DataSourceExtras
 import com.databricks.labs.gbx.vectorx.ds.ogr.OGR_DataSource
+import org.apache.spark.sql.connector.catalog.Table
+import org.apache.spark.sql.connector.expressions.Transform
+import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
+
 
 //noinspection ScalaUnusedSymbol
-class ShapeFile_DataSource extends OGR_DataSource {
+class ShapeFile_DataSource extends OGR_DataSource with DataSourceExtras{
+
+    override def dsExtraMap(checkMap: Map[String, String] = Map.empty): Map[String, String] = Map(
+        "driverName" -> "ESRI Shapefile"
+    )
 
     override def shortName(): String = "shapefile"
 
+    override def inferSchema(options: CaseInsensitiveStringMap): StructType = {
+        super.inferSchema(extraCaseInsensitiveStringMap(options))
+    }
+
+    override def getTable(schema: StructType, partitions: Array[Transform], properties: java.util.Map[String, String]): Table = {
+        super.getTable(schema, partitions, extraJavaUtilMap(properties))
+    }
 }
diff --git a/src/test/resources/binary/gdb/bridges.gdb.zip b/src/test/resources/binary/gdb/bridges.gdb.zip
diff --git a/src/test/resources/binary/geoparquet/data-multipolygon-encoding_wkb.parquet b/src/test/resources/binary/geoparquet/data-multipolygon-encoding_wkb.parquet
diff --git a/src/test/resources/binary/gpkg/util_wastewater_discharge.gpkg b/src/test/resources/binary/gpkg/util_wastewater_discharge.gpkg
diff --git a/src/test/resources/binary/shapefile/zip/tl_rd22_13037_addrfeat.zip b/src/test/resources/binary/shapefile/zip/tl_rd22_13037_addrfeat.zip
diff --git a/src/test/scala/com/databricks/labs/gbx/vectorx/ds/OGR_DataSourceTest.scala b/src/test/scala/com/databricks/labs/gbx/vectorx/ds/OGR_DataSourceTest.scala
@@ -3,19 +3,23 @@ package com.databricks.labs.gbx.vectorx.ds
 import org.apache.spark.sql.catalyst.plans.PlanTest
 import org.apache.spark.sql.functions._
 import org.apache.spark.sql.test.SilentSparkSession
+import org.apache.spark.sql.util.CaseInsensitiveStringMap
 import org.scalatest.matchers.must.Matchers.be
 import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
 
+
 class OGR_DataSourceTest extends PlanTest with SilentSparkSession {
 
-    test("GDAL Data Source must read tif files") {
+    test("GDAL Data Source must read OGR files") {
         val sp = spark
         import com.databricks.labs.gbx.udfs._
         import sp.implicits._
 
+        // shapefile
+
         val shpPath = this.getClass.getResource("/binary/elevation/sd46_dtm_breakline.shp").toString.replace("file:", "")
 
-        val res = spark.read
+        val res_sp1 = spark.read
             .format("ogr")
             .option("chunkSize", "100")
             .load(shpPath)
@@ -24,14 +28,61 @@ class OGR_DataSourceTest extends PlanTest with SilentSparkSession {
             .as[Double]
             .collect()
 
-        res.foreach(v => v should be >= 0.0)
+        res_sp1.foreach(v => v should be >= 0.0)
 
-        val res2 = spark.read
+        val res_sp2 = spark.read
             .format("shapefile")
             .option("chunkSize", "100")
             .load(shpPath)
 
-        res2.count() should be > 0L
+        res_sp2.count() should be > 0L
+
+        // shapefile zip
+
+        val shpZipPath = this.getClass.getResource("/binary/shapefile/zip/tl_rd22_13037_addrfeat.zip").toString.replace("file:", "")
+
+        val res_spz = spark.read
+            .format("shapefile")
+            .load(shpZipPath)
+
+        res_spz.count() should be > 0L
+
+        // file_gdb
+
+        val gdbZipPath = this.getClass.getResource("/binary/gdb/bridges.gdb.zip").toString.replace("file:", "")
+
+        val res_gdb = spark.read
+            .format("file_gdb")
+            .load(gdbZipPath)
+
+        res_gdb.count() should be > 0L
+
+        // geojson
+
+        val gjPath = this.getClass.getResource("/NYC_Taxi_Zones.geojson").toString.replace("file:", "")
+
+        val res_gj = spark.read
+            .format("geojson")
+            .option("multi", "false")
+            .load(gjPath)
+
+        res_gj.count() shouldEqual 1L // single geom
+
+        val res_gj1 = spark.read
+            .format("geojson")
+            .load(gjPath)
+
+        res_gj1.count() should be > 1L // newline geoms
+
+        // gpkg (zip fails)
+
+        val gpkgPath = this.getClass.getResource("/binary/gpkg/util_wastewater_discharge.gpkg").toString.replace("file:", "")
+
+        val res_gpkg = spark.read
+            .format("ogr_gpkg")
+            .load(gpkgPath)
+
+        res_gpkg.count() should be > 0L
 
     }
 
