jiayuasu commented on code in PR #2846: URL: https://github.com/apache/sedona/pull/2846#discussion_r3122127884
########## spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/io/geotiffmetadata/GeoTiffMetadataPartitionReader.scala: ########## @@ -0,0 +1,396 @@ +/* + * 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.geotiffmetadata + +import org.apache.hadoop.conf.Configuration +import org.apache.hadoop.fs.Path +import org.apache.sedona.common.raster.RasterAccessors +import org.apache.sedona.common.raster.RasterBandAccessors +import org.apache.sedona.common.raster.inputstream.HadoopImageInputStream +import org.apache.sedona.common.utils.RasterUtils +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow +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 org.geotools.coverage.grid.GridCoverage2D +import org.geotools.gce.geotiff.GeoTiffReader +import org.geotools.referencing.crs.DefaultEngineeringCRS + +import java.net.URI +import scala.collection.mutable +import scala.util.Try + +class GeoTiffMetadataPartitionReader( + 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 imageStream = new HadoopImageInputStream(path, configuration) + var reader: GeoTiffReader = null + var raster: GridCoverage2D = null + try { + reader = new GeoTiffReader( + imageStream, + new org.geotools.util.factory.Hints( + org.geotools.util.factory.Hints.FORCE_LONGITUDE_FIRST_AXIS_ORDER, + java.lang.Boolean.TRUE)) + + // Extract TIFF IIO metadata BEFORE read() to avoid stream state issues + val isTiled = GeoTiffMetadataPartitionReader.hasTiffTag(reader, 322) + val photometric = GeoTiffMetadataPartitionReader.extractPhotometricInterpretation(reader) + val tiffMetadata = GeoTiffMetadataPartitionReader.extractMetadata(reader) + val compression = GeoTiffMetadataPartitionReader.extractCompression(reader) + + raster = reader.read(null) + + lazy val width = RasterAccessors.getWidth(raster) + lazy val height = RasterAccessors.getHeight(raster) + lazy val numBands = RasterAccessors.numBands(raster) + lazy val srid = RasterAccessors.srid(raster) + + lazy val crsStr = Try { + val crs = raster.getCoordinateReferenceSystem + if (crs == null || crs.isInstanceOf[DefaultEngineeringCRS]) null + else crs.toWKT + }.getOrElse(null) + + lazy val affine = RasterUtils.getGDALAffineTransform(raster) + lazy val geoTransformRow = new GenericInternalRow( + Array[Any]( + affine.getTranslateX, + affine.getTranslateY, + affine.getScaleX, + affine.getScaleY, + affine.getShearX, + affine.getShearY)) + + lazy val env = raster.getEnvelope2D + lazy val cornerCoordinatesRow = new GenericInternalRow( + Array[Any](env.getMinX, env.getMinY, env.getMaxX, env.getMaxY)) + + lazy val image = raster.getRenderedImage + lazy val tileWidth = image.getTileWidth + lazy val tileHeight = image.getTileHeight + + lazy val bandsArray = GeoTiffMetadataPartitionReader + .buildBandsArray(raster, numBands, tileWidth, tileHeight, photometric) + lazy val overviewsArray = + GeoTiffMetadataPartitionReader.buildOverviewsArray(reader, width, height) + lazy val metadataMap = GeoTiffMetadataPartitionReader.buildMetadataMap(tiffMetadata) + + val fields = readDataSchema.fieldNames.map { + case "path" => UTF8String.fromString(path.toString) + case "driver" => UTF8String.fromString("GTiff") + case "fileSize" => partition.fileSize: Any + case "width" => width: Any + case "height" => height: Any + case "numBands" => numBands: Any + case "srid" => srid: Any + case "crs" => + if (crsStr != null) UTF8String.fromString(crsStr) else null + case "geoTransform" => geoTransformRow + case "cornerCoordinates" => cornerCoordinatesRow + case "bands" => bandsArray + case "overviews" => overviewsArray + case "metadata" => metadataMap + case "isTiled" => isTiled: Any + case "compression" => + if (compression != null) UTF8String.fromString(compression) else null + case other => + throw new IllegalArgumentException(s"Unsupported field name: $other") + } + + new GenericInternalRow(fields) + } finally { + if (raster != null) raster.dispose(true) + if (reader != null) reader.dispose() + imageStream.close() + } + } +} + +object GeoTiffMetadataPartitionReader { + + // TIFF Photometric Interpretation values + private val PHOTOMETRIC_MIN_IS_WHITE = 0 + private val PHOTOMETRIC_MIN_IS_BLACK = 1 + private val PHOTOMETRIC_RGB = 2 + private val PHOTOMETRIC_PALETTE = 3 + + def buildBandsArray( + raster: GridCoverage2D, + numBands: Int, + tileWidth: Int, + tileHeight: Int, + photometric: Int): GenericArrayData = { + val bands = (1 to numBands).map { i => + val dataType = Try(RasterBandAccessors.getBandType(raster, i)).getOrElse(null) + val noDataValue = Try(RasterBandAccessors.getBandNoDataValue(raster, i)).getOrElse(null) + val description = Try { + val desc = raster.getSampleDimension(i - 1).getDescription + if (desc != null) desc.toString(java.util.Locale.ROOT) else null + }.getOrElse(null) + val unit = Try { + val units = raster.getSampleDimension(i - 1).getUnits + if (units != null) units.toString else null + }.getOrElse(null) + val colorInterp = resolveColorInterpretation(photometric, i, numBands) + + new GenericInternalRow( + Array[Any]( + i, + if (dataType != null) UTF8String.fromString(dataType) else null, + if (colorInterp != null) UTF8String.fromString(colorInterp) else null, + if (noDataValue != null) noDataValue.doubleValue() else null, + tileWidth, + tileHeight, + if (description != null) UTF8String.fromString(description) else null, + if (unit != null) UTF8String.fromString(unit) else null)) + }.toArray + new GenericArrayData(bands) + } + + private def resolveColorInterpretation(photometric: Int, band: Int, numBands: Int): String = { + photometric match { + case PHOTOMETRIC_MIN_IS_WHITE | PHOTOMETRIC_MIN_IS_BLACK => + if (numBands == 1) "Gray" + else s"Gray$band" + case PHOTOMETRIC_RGB => + band match { + case 1 => "Red" + case 2 => "Green" + case 3 => "Blue" + case 4 => "Alpha" + case _ => "Undefined" + } + case PHOTOMETRIC_PALETTE => "Palette" + case _ => "Undefined" + } + } + + def buildOverviewsArray( + reader: GeoTiffReader, + fullWidth: Int, + fullHeight: Int): GenericArrayData = { + try { + val layout = reader.getDatasetLayout + if (layout == null) return new GenericArrayData(Array.empty[InternalRow]) + + val numOverviews = layout.getNumInternalOverviews + if (numOverviews <= 0) return new GenericArrayData(Array.empty[InternalRow]) + + val resolutionLevels = reader.getResolutionLevels + if (resolutionLevels == null || resolutionLevels.length <= 1) + return new GenericArrayData(Array.empty[InternalRow]) + + val count = Math.min(numOverviews, resolutionLevels.length - 1) + val fullResX = resolutionLevels(0)(0) + val fullResY = resolutionLevels(0)(1) + val overviews = (1 to count).map { level => + val ovrResX = resolutionLevels(level)(0) + val ovrResY = resolutionLevels(level)(1) + new GenericInternalRow( + Array[Any]( + level, + Math.round(fullWidth.toDouble * fullResX / ovrResX).toInt, + Math.round(fullHeight.toDouble * fullResY / ovrResY).toInt)) + }.toArray + new GenericArrayData(overviews) + } catch { + case _: Exception => new GenericArrayData(Array.empty[InternalRow]) + } + } + + def buildMetadataMap( + metadata: Map[String, String]): org.apache.spark.sql.catalyst.util.MapData = { + if (metadata.isEmpty) return null + org.apache.spark.sql.catalyst.util.ArrayBasedMapData( + metadata.keys.map(UTF8String.fromString).toArray, + metadata.values.map(UTF8String.fromString).toArray) Review Comment: Fixed in 6d408e2e58. Rebuilt key/value arrays from a single pass over `(k, v)` entries. ########## spark/common/src/main/scala/org/apache/spark/sql/sedona_sql/io/geotiffmetadata/GeoTiffMetadataTable.scala: ########## @@ -0,0 +1,113 @@ +/* + * 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.geotiffmetadata + +import org.apache.hadoop.fs.FileStatus +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.connector.catalog.SupportsRead +import org.apache.spark.sql.connector.catalog.TableCapability +import org.apache.spark.sql.connector.read.ScanBuilder +import org.apache.spark.sql.connector.write.LogicalWriteInfo +import org.apache.spark.sql.connector.write.WriteBuilder +import org.apache.spark.sql.execution.datasources.FileFormat +import org.apache.spark.sql.execution.datasources.v2.FileTable +import org.apache.spark.sql.types._ +import org.apache.spark.sql.util.CaseInsensitiveStringMap + +import java.util.{Set => JSet} + +case class GeoTiffMetadataTable( + name: String, + sparkSession: SparkSession, + options: CaseInsensitiveStringMap, + paths: Seq[String], + userSpecifiedSchema: Option[StructType], + fallbackFileFormat: Class[_ <: FileFormat]) + extends FileTable(sparkSession, options, paths, userSpecifiedSchema) + with SupportsRead { + + override def inferSchema(files: Seq[FileStatus]): Option[StructType] = + Some(userSpecifiedSchema.getOrElse(GeoTiffMetadataTable.SCHEMA)) + + override def formatName: String = "GeoTiffMetadata" + + override def capabilities(): JSet[TableCapability] = + java.util.EnumSet.of(TableCapability.BATCH_READ) + + override def newScanBuilder(options: CaseInsensitiveStringMap): ScanBuilder = { + GeoTiffMetadataScanBuilder(sparkSession, fileIndex, schema, dataSchema, options) + } + + override def newWriteBuilder(info: LogicalWriteInfo): WriteBuilder = + throw new UnsupportedOperationException("GeoTiffMetadata is a read-only data source") +} + +object GeoTiffMetadataTable { + + val GEO_TRANSFORM_TYPE: StructType = StructType( + Seq( + StructField("upperLeftX", DoubleType, nullable = false), + StructField("upperLeftY", DoubleType, nullable = false), + StructField("scaleX", DoubleType, nullable = false), + StructField("scaleY", DoubleType, nullable = false), + StructField("skewX", DoubleType, nullable = false), + StructField("skewY", DoubleType, nullable = false))) + + val CORNER_COORDINATES_TYPE: StructType = StructType( + Seq( + StructField("minX", DoubleType, nullable = false), + StructField("minY", DoubleType, nullable = false), + StructField("maxX", DoubleType, nullable = false), + StructField("maxY", DoubleType, nullable = false))) + + val BAND_TYPE: StructType = StructType( + Seq( + StructField("band", IntegerType, nullable = false), + StructField("dataType", StringType, nullable = true), + StructField("colorInterpretation", StringType, nullable = true), + StructField("noDataValue", DoubleType, nullable = true), + StructField("blockWidth", IntegerType, nullable = false), + StructField("blockHeight", IntegerType, nullable = false), + StructField("description", StringType, nullable = true), + StructField("unit", StringType, nullable = true))) + + val OVERVIEW_TYPE: StructType = StructType( + Seq( + StructField("level", IntegerType, nullable = false), + StructField("width", IntegerType, nullable = false), + StructField("height", IntegerType, nullable = false))) + + val SCHEMA: StructType = StructType( + Seq( + StructField("path", StringType, nullable = false), + StructField("driver", StringType, nullable = false), + StructField("fileSize", LongType, nullable = true), Review Comment: Fixed in 6d408e2e58. Marked `fileSize` as `nullable = false`. -- 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]
