Copilot commented on code in PR #2826: URL: https://github.com/apache/sedona/pull/2826#discussion_r3045682533
########## spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/io/sedonainfo/SedonaInfoPartitionReader.scala: ########## @@ -0,0 +1,151 @@ +/* + * 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.spark.sql.sedona_sql.io.sedonainfo + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +import org.apache.spark.sql.catalyst.util.ArrayData +import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.connector.read.PartitionReader +import org.apache.spark.sql.execution.datasources.PartitionedFile +import org.apache.spark.sql.types.StructType +import org.apache.spark.unsafe.types.UTF8String + +import java.net.URI + +/** + * Reads raster file metadata by delegating to format-specific [[RasterFileMetadataExtractor]] + * implementations. Produces one [[InternalRow]] per file matching the readDataSchema. + */ +class SedonaInfoPartitionReader( + configuration: Configuration, + partitionedFiles: Array[PartitionedFile], + readDataSchema: StructType) + extends PartitionReader[InternalRow] { + + private var currentFileIndex = 0 + private var currentRow: InternalRow = _ + + override def next(): Boolean = { + if (currentFileIndex < partitionedFiles.length) { + currentRow = readFileMetadata(partitionedFiles(currentFileIndex)) + currentFileIndex += 1 + true + } else { + false + } + } + + override def get(): InternalRow = currentRow + + override def close(): Unit = {} + + private def readFileMetadata(partition: PartitionedFile): InternalRow = { + val path = new Path(new URI(partition.filePath.toString())) + val extractor = SedonaInfoPartitionReader.findExtractor(path) + val requiredFields = readDataSchema.fieldNames.toSet + val meta = extractor.extract(path, partition.fileSize, configuration, requiredFields) + SedonaInfoPartitionReader.toInternalRow(meta, readDataSchema) + } +} + +object SedonaInfoPartitionReader { + + /** Registered metadata extractors. Add new format extractors here. */ + private val extractors: Seq[RasterFileMetadataExtractor] = Seq(GeoTiffMetadataExtractor) + + def findExtractor(path: Path): RasterFileMetadataExtractor = { + extractors + .find(_.canHandle(path)) + .getOrElse( + throw new UnsupportedOperationException( + s"No metadata extractor found for file: ${path.getName}. " + + s"Supported formats: ${extractors.map(_.driver).mkString(", ")}")) + } + + def toInternalRow(meta: RasterFileMetadata, readDataSchema: StructType): InternalRow = { + val gt = meta.geoTransform + val geoTransformRow = new GenericInternalRow( + Array[Any](gt.upperLeftX, gt.upperLeftY, gt.scaleX, gt.scaleY, gt.skewX, gt.skewY)) + + val cc = meta.cornerCoordinates + val cornerCoordinatesRow = + new GenericInternalRow(Array[Any](cc.minX, cc.minY, cc.maxX, cc.maxY)) + + lazy val bandsArray: ArrayData = { + val bands = meta.bands.map { b => + new GenericInternalRow( + Array[Any]( + b.band, + if (b.dataType != null) UTF8String.fromString(b.dataType) else null, + if (b.colorInterpretation != null) UTF8String.fromString(b.colorInterpretation) + else null, + if (b.noDataValue != null) b.noDataValue.doubleValue() else null, + b.blockWidth, + b.blockHeight, + if (b.description != null) UTF8String.fromString(b.description) else null, + if (b.unit != null) UTF8String.fromString(b.unit) else null)) + }.toArray + new GenericArrayData(bands) + } + + lazy val overviewsArray: ArrayData = { + val overviews = meta.overviews.map { o => + new GenericInternalRow(Array[Any](o.level, o.width, o.height)) + }.toArray + new GenericArrayData(overviews) + } + + lazy val metadataMap: Any = { + if (meta.metadata.isEmpty) null + else { + val entries = meta.metadata.toSeq + org.apache.spark.sql.catalyst.util.ArrayBasedMapData( + entries.map { case (k, _) => UTF8String.fromString(k) }.toArray, + entries.map { case (_, v) => UTF8String.fromString(v) }.toArray) + } Review Comment: `metadataMap` is set to `null` when `meta.metadata.isEmpty`, while other collection fields like `bands`/`overviews` return empty arrays. Returning `null` makes common Spark operations (e.g., `size(metadata)`) yield `null` instead of `0` and forces extra null-handling in user queries. Consider returning an empty map value (empty key/value arrays) instead of `null` for the no-metadata case for consistency and easier querying. ```suggestion val entries = meta.metadata.toSeq org.apache.spark.sql.catalyst.util.ArrayBasedMapData( entries.map { case (k, _) => UTF8String.fromString(k) }.toArray, entries.map { case (_, v) => UTF8String.fromString(v) }.toArray) ``` ########## spark/common/src/test/scala/org/apache/sedona/sql/geotiffMetadataTest.scala: ########## @@ -0,0 +1,222 @@ +/* + * 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.sql + +import org.apache.commons.io.FileUtils +import org.junit.Assert.assertEquals +import org.scalatest.BeforeAndAfterAll + +import java.io.File +import java.nio.file.Files + +class geotiffMetadataTest extends TestBaseScala with BeforeAndAfterAll { + + val rasterDir: String = resourceFolder + "raster/" + val singleFileLocation: String = resourceFolder + "raster/test1.tiff" + val tempDir: String = + Files.createTempDirectory("sedona_sedonainfo_test_").toFile.getAbsolutePath + + override def afterAll(): Unit = { + FileUtils.deleteDirectory(new File(tempDir)) + super.afterAll() + } + + describe("SedonaInfo data source") { + + it("should read test1.tiff with exact metadata values") { + val df = sparkSession.read.format("sedonainfo").load(singleFileLocation) + assert(df.count() == 1) + + val row = df.first() + assert(row.getAs[String]("path").endsWith("test1.tiff")) + assertEquals("GTiff", row.getAs[String]("driver")) + assertEquals(174803L, row.getAs[Long]("fileSize")) + assertEquals(512, row.getAs[Int]("width")) + assertEquals(517, row.getAs[Int]("height")) + assertEquals(1, row.getAs[Int]("numBands")) + assertEquals(3857, row.getAs[Int]("srid")) + assert(row.getAs[String]("crs").contains("EPSG")) + // test1.tiff has TileWidth/TileLength TIFF tags (internally tiled) + assertEquals(true, row.getAs[Boolean]("isTiled")) + } + + it("should return exact geoTransform for test1.tiff") { + val row = sparkSession.read + .format("sedonainfo") + .load(singleFileLocation) + .selectExpr( + "geoTransform.upperLeftX", + "geoTransform.upperLeftY", + "geoTransform.scaleX", + "geoTransform.scaleY", + "geoTransform.skewX", + "geoTransform.skewY") + .first() + assertEquals(-1.3095817809482181e7, row.getAs[Double]("upperLeftX"), 0.01) + assertEquals(4021262.7487925636, row.getAs[Double]("upperLeftY"), 0.01) + assertEquals(72.32861272132695, row.getAs[Double]("scaleX"), 1e-10) + assertEquals(-72.32861272132695, row.getAs[Double]("scaleY"), 1e-10) + assertEquals(0.0, row.getAs[Double]("skewX"), 1e-15) + assertEquals(0.0, row.getAs[Double]("skewY"), 1e-15) + } + + it("should return exact cornerCoordinates for test1.tiff") { + val row = sparkSession.read + .format("sedonainfo") + .load(singleFileLocation) + .selectExpr( + "cornerCoordinates.minX", + "cornerCoordinates.minY", + "cornerCoordinates.maxX", + "cornerCoordinates.maxY") + .first() + assertEquals(-1.3095817809482181e7, row.getAs[Double]("minX"), 0.01) + assertEquals(3983868.8560156375, row.getAs[Double]("minY"), 0.01) + assertEquals(-1.3058785559768861e7, row.getAs[Double]("maxX"), 0.01) + assertEquals(4021262.7487925636, row.getAs[Double]("maxY"), 0.01) + } + + it("should return exact band metadata for test1.tiff") { + val row = sparkSession.read + .format("sedonainfo") + .load(singleFileLocation) + .selectExpr("explode(bands) as b") + .selectExpr( + "b.band", + "b.dataType", + "b.colorInterpretation", + "b.noDataValue", + "b.blockWidth", + "b.blockHeight", + "b.description", + "b.unit") + .first() + assertEquals(1, row.getAs[Int]("band")) + assertEquals("UNSIGNED_8BITS", row.getAs[String]("dataType")) + assertEquals("Gray", row.getAs[String]("colorInterpretation")) + assert(row.isNullAt(row.fieldIndex("noDataValue"))) + assertEquals(256, row.getAs[Int]("blockWidth")) + assertEquals(256, row.getAs[Int]("blockHeight")) + assertEquals("GRAY_INDEX", row.getAs[String]("description")) + assert(row.isNullAt(row.fieldIndex("unit"))) + } + + it("should return empty overviews for non-COG test1.tiff") { + // test1.tiff has only 1 IFD (no internal overviews) + val row = sparkSession.read + .format("sedonainfo") + .load(singleFileLocation) + .selectExpr("size(overviews) as overviewCount") + .first() + assertEquals(0, row.getAs[Int]("overviewCount")) + } + + it("should cross-validate against raster data source") { + val metaRow = sparkSession.read.format("sedonainfo").load(singleFileLocation).first() + val rasterRow = sparkSession.read + .format("raster") + .option("retile", "false") + .load(singleFileLocation) + .selectExpr( + "RS_Width(rast) as width", + "RS_Height(rast) as height", + "RS_NumBands(rast) as numBands", + "RS_SRID(rast) as srid") + .first() + assertEquals(metaRow.getAs[Int]("width"), rasterRow.getAs[Int]("width")) + assertEquals(metaRow.getAs[Int]("height"), rasterRow.getAs[Int]("height")) + assertEquals(metaRow.getAs[Int]("numBands"), rasterRow.getAs[Int]("numBands")) + assertEquals(metaRow.getAs[Int]("srid"), rasterRow.getAs[Int]("srid")) + } + + it("should read multiple files via glob") { + val df = sparkSession.read.format("sedonainfo").load(rasterDir + "*.tiff") + // 7 .tiff files in the raster directory (excludes test3.tif) + assertEquals(7L, df.count()) + } + + it("should read files from directory with trailing slash") { + val df = sparkSession.read.format("sedonainfo").load(rasterDir) + // Recursive lookup finds all .tif/.tiff files including subdirectories + assertEquals(9L, df.count()) + } + + it("should support LIMIT pushdown") { + val df = sparkSession.read.format("sedonainfo").load(rasterDir).limit(2) + assertEquals(2L, df.count()) + } Review Comment: The test named "should support LIMIT pushdown" only asserts `df.limit(2).count() == 2`, which would pass even if LIMIT is not pushed down (Spark can apply the limit after reading all rows). To actually validate pushdown, consider asserting on the optimized/executed plan (e.g., presence of the scan’s pushed limit) or otherwise verifying that only 2 files were planned/read. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
