GEOMESA-3577 FSDS - Add z-value column for partition pruning (#3543)

Author: elahrvivaz

docs/user/filesystem/configuration.rst

@@ -76,6 +76,13 @@ are:
 * ``GeoParquetNative`` - This schema uses `GeoParquet 1.1.0 <https://geoparquet.org/releases/v1.1.0/>`__ with geometries
   encoded "natively". This format doesn't require special libraries to read, but isn't as widely supported as WKB.
 
+geomesa.parquet.z-value-column
+++++++++++++++++++++++++++++++
+
+This property can be used to skip writing an additional Z-value column for geometry-type columns when using Parquet, by setting
+it to ``false``. By default, each geometry will include an int-64-type column that includes the index value from a Z2 or XZ2
+space-filling-curve (as appropriate). This can be used to accelerate queries through partition pruning.
+
 AWS S3 Configuration
 --------------------
 

geomesa-convert/geomesa-convert-parquet/src/main/scala/org/locationtech/geomesa/convert/parquet/AvroReadSupport.scala

@@ -21,8 +21,9 @@ import org.apache.parquet.schema.Type.Repetition
 import org.apache.parquet.schema._
 import org.locationtech.geomesa.convert.parquet.AvroReadSupport.AvroRecordMaterializer
 import org.locationtech.geomesa.curve.BinnedTime
-import org.locationtech.geomesa.fs.storage.parquet.io.SimpleFeatureReadSupport._
-import org.locationtech.geomesa.fs.storage.parquet.io.{SimpleFeatureParquetSchema, SimpleFeatureReadSupport}
+import org.locationtech.geomesa.fs.storage.parquet.io.SimpleFeatureParquetSchema
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.SimpleFeatureReadSupport
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.SimpleFeatureReadSupport._
 
 import java.util.{Collections, Date}
 

geomesa-convert/geomesa-convert-parquet/src/main/scala/org/locationtech/geomesa/convert/parquet/ParquetConverterFactory.scala

@@ -27,7 +27,7 @@ import org.locationtech.geomesa.convert2.TypeInference.{FunctionTransform, Infer
 import org.locationtech.geomesa.convert2.transforms.Expression
 import org.locationtech.geomesa.convert2.{AbstractConverterFactory, TypeInference}
 import org.locationtech.geomesa.fs.storage.parquet.io.GeoParquetMetadata.{ColumnMetadata, GeoParquetColumnEncoding, GeoParquetColumnType}
-import org.locationtech.geomesa.fs.storage.parquet.io.GeometrySchema.GeometryEncoding.GeoParquetNative
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.GeometrySchema.GeometryEncoding.GeoParquetNative
 import org.locationtech.geomesa.fs.storage.parquet.io.{GeoParquetMetadata, SimpleFeatureParquetSchema}
 import org.locationtech.geomesa.utils.geotools.ObjectType.ObjectType
 import org.locationtech.geomesa.utils.geotools.{ObjectType, SimpleFeatureTypes}

geomesa-convert/geomesa-convert-parquet/src/main/scala/org/locationtech/geomesa/convert/parquet/ParquetFunctionFactory.scala

@@ -60,7 +60,7 @@ class ParquetFunctionFactory extends TransformerFunctionFactory {
   abstract class ParquetGeometryFn[T <: Geometry, U](name: String, path: Option[AvroPath])
       extends NamedTransformerFunction(Seq(name), pure = true) {
 
-    import org.locationtech.geomesa.fs.storage.parquet.io.GeometrySchema.{GeometryColumnX, GeometryColumnY}
+    import org.locationtech.geomesa.fs.storage.parquet.io.geometry.GeometrySchema.{GeometryColumnX, GeometryColumnY}
 
     override def apply(args: Array[AnyRef]): AnyRef = {
       val attribute = path match {

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-jobs/src/main/scala/org/locationtech/geomesa/fs/storage/jobs/parquet/ParquetSimpleFeatureInputFormat.scala

@@ -19,7 +19,7 @@ import org.geotools.api.filter.Filter
 import org.locationtech.geomesa.features.TransformSimpleFeature
 import org.locationtech.geomesa.fs.storage.jobs.StorageConfiguration
 import org.locationtech.geomesa.fs.storage.jobs.parquet.ParquetSimpleFeatureInputFormat.{ParquetSimpleFeatureInputFormatBase, ParquetSimpleFeatureRecordReaderBase, ParquetSimpleFeatureTransformRecordReaderBase}
-import org.locationtech.geomesa.fs.storage.parquet.io.SimpleFeatureReadSupport
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.SimpleFeatureReadSupport
 import org.locationtech.geomesa.fs.storage.parquet.{ReadFilter, ReadSchema}
 import org.locationtech.geomesa.index.planning.QueryRunner
 

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-jobs/src/main/scala/org/locationtech/geomesa/fs/storage/jobs/parquet/ParquetStorageConfiguration.scala

@@ -14,7 +14,7 @@ import org.apache.parquet.hadoop.metadata.CompressionCodecName
 import org.apache.parquet.hadoop.{ParquetInputFormat, ParquetOutputFormat}
 import org.geotools.api.feature.simple.SimpleFeatureType
 import org.locationtech.geomesa.fs.storage.jobs.StorageConfiguration
-import org.locationtech.geomesa.fs.storage.parquet.io.{SimpleFeatureReadSupport, SimpleFeatureWriteSupport}
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.{SimpleFeatureReadSupport, SimpleFeatureWriteSupport}
 
 trait ParquetStorageConfiguration extends StorageConfiguration with LazyLogging {
   override def configureOutput(sft: SimpleFeatureType, job: Job): Unit = {

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-parquet-io/src/main/scala/org/locationtech/geomesa/fs/storage/parquet/io/GeoParquetMetadata.scala

@@ -13,7 +13,7 @@ import org.geotools.api.feature.`type`.GeometryDescriptor
 import org.geotools.api.feature.simple.SimpleFeature
 import org.locationtech.geomesa.fs.storage.core.observer.FileSystemObserver
 import org.locationtech.geomesa.fs.storage.parquet.io.GeoParquetMetadata.ColumnMetadata
-import org.locationtech.geomesa.fs.storage.parquet.io.GeometrySchema.GeometryEncoding
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.GeometrySchema.GeometryEncoding
 import org.locationtech.geomesa.utils.geotools.ObjectType
 import org.locationtech.geomesa.utils.geotools.ObjectType.ObjectType
 import org.locationtech.geomesa.utils.text.StringSerialization

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-parquet-io/src/main/scala/org/locationtech/geomesa/fs/storage/parquet/io/ParquetFileSystemReader.scala

@@ -19,6 +19,7 @@ import org.locationtech.geomesa.features.TransformSimpleFeature
 import org.locationtech.geomesa.fs.storage.core.FileSystemContext
 import org.locationtech.geomesa.fs.storage.core.FileSystemStorage.FileSystemPathReader
 import org.locationtech.geomesa.fs.storage.core.fs.{LocalObjectStore, ObjectStore, S3ObjectStore}
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.SimpleFeatureReadSupport
 import org.locationtech.geomesa.fs.storage.parquet.io.s3.S3InputFile
 import org.locationtech.geomesa.utils.collection.CloseableIterator
 import org.locationtech.geomesa.utils.geotools.Transform.Transforms

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-parquet-io/src/main/scala/org/locationtech/geomesa/fs/storage/parquet/io/ParquetFileSystemWriter.scala

@@ -21,6 +21,7 @@ import org.locationtech.geomesa.fs.storage.core.FileSystemStorage.FileSystemWrit
 import org.locationtech.geomesa.fs.storage.core.fs.{LocalObjectStore, ObjectStore, S3ObjectStore}
 import org.locationtech.geomesa.fs.storage.core.observer.FileSystemObserver
 import org.locationtech.geomesa.fs.storage.core.observer.FileSystemObserverFactory.NoOpObserver
+import org.locationtech.geomesa.fs.storage.parquet.io.rw.SimpleFeatureWriteSupport
 import org.locationtech.geomesa.fs.storage.parquet.io.s3.S3OutputFile
 import org.locationtech.geomesa.utils.io.CloseQuietly
 

geomesa-fs/geomesa-fs-storage/geomesa-fs-storage-parquet-io/src/main/scala/org/locationtech/geomesa/fs/storage/parquet/io/SimpleFeatureParquetSchema.scala

@@ -10,18 +10,20 @@
 package org.locationtech.geomesa.fs.storage.parquet.io
 
 import com.typesafe.scalalogging.LazyLogging
-import org.apache.hadoop.conf.Configuration
 import org.apache.iceberg.Schema
 import org.apache.iceberg.types.Types._
-import org.apache.parquet.conf.{HadoopParquetConfiguration, ParquetConfiguration, PlainParquetConfiguration}
+import org.apache.parquet.conf.{ParquetConfiguration, PlainParquetConfiguration}
 import org.apache.parquet.hadoop.api.InitContext
 import org.apache.parquet.hadoop.metadata.FileMetaData
 import org.apache.parquet.schema.LogicalTypeAnnotation.TimeUnit
 import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
 import org.apache.parquet.schema.Type.Repetition
 import org.apache.parquet.schema._
 import org.geotools.api.feature.simple.SimpleFeatureType
-import org.locationtech.geomesa.fs.storage.parquet.io.GeometrySchema.{BoundingBoxField, BoundingBoxes, GeometryEncoding}
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.BoundingBoxes.BoundingBoxField
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.GeometrySchema.GeometryEncoding
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.ZValues.ZValueField
+import org.locationtech.geomesa.fs.storage.parquet.io.geometry.{BoundingBoxes, GeometrySchema, ZValues}
 import org.locationtech.geomesa.utils.geotools.ObjectType.ObjectType
 import org.locationtech.geomesa.utils.geotools.{ObjectType, SimpleFeatureTypes}
 import org.locationtech.geomesa.utils.text.StringSerialization
@@ -35,6 +37,7 @@ import java.util.concurrent.atomic.AtomicInteger
  * @param encodings type encoding
  * @param hasVisibilities whether the schema encodes visibilities, or not
  * @param bboxes fields with bounding boxes
+ * @param zValues fields with z-value columns
  * @param metadata file metadata
  * @param schema parquet message schema
  */
@@ -43,6 +46,7 @@ case class SimpleFeatureParquetSchema(
     encodings: Encodings,
     hasVisibilities: Boolean,
     bboxes: BoundingBoxes,
+    zValues: ZValues,
     metadata: java.util.Map[String, String],
     schema: MessageType) {
 
@@ -71,6 +75,7 @@ object SimpleFeatureParquetSchema extends LazyLogging {
   val SftReadSpecKey        = "geomesa.fs.sft.read.spec"
   val GeometryEncodingKey   = "geomesa.parquet.geometries"
   val BBoxEncodingKey       = "geomesa.parquet.bounding-boxes"
+  val ZValueColumKey        = "geomesa.parquet.z-value-column"
   val VisibilityEncodingKey = "geomesa.fs.visibilities"
   val PartitionKey          = "geomesa.fs.partition"
 
@@ -148,6 +153,7 @@ object SimpleFeatureParquetSchema extends LazyLogging {
       // only include bboxes if they help with push-down filters
       val bboxes =
         if (Option(conf.get(BBoxEncodingKey)).forall(_.toBoolean)) { BoundingBoxes(sft, geometries) } else { BoundingBoxes(Seq.empty) }
+      val zCols = if (Option(conf.get(ZValueColumKey)).forall(_.toBoolean)) { ZValues(sft) } else { ZValues(Seq.empty) }
       val visibilities =
         Option(sft.getUserData.get(VisibilityEncodingKey)).orElse(Option(conf.get(VisibilityEncodingKey))).forall(_.toString.toBoolean)
       val metadata = Map(
@@ -156,8 +162,8 @@ object SimpleFeatureParquetSchema extends LazyLogging {
         GeometryEncodingKey -> geometries.toString,
       ) ++ Option(conf.get(PartitionKey)).map(PartitionKey -> _)
       val encodings = Encodings(geometries)
-      val messageType = schema(sft, None, encodings, bboxes, visibilities)
-      SimpleFeatureParquetSchema(sft, encodings, visibilities, bboxes, metadata.asJava, messageType)
+      val messageType = schema(sft, encodings, bboxes, zCols, visibilities).toMessageType(sft)
+      SimpleFeatureParquetSchema(sft, encodings, visibilities, bboxes, zCols, metadata.asJava, messageType)
     }
   }
 
@@ -192,10 +198,11 @@ object SimpleFeatureParquetSchema extends LazyLogging {
         throw new UnsupportedOperationException("GeoMesaV0/GeoMesaV1 encoding is no longer supported")
       }
       val encodings = Encodings(geometries)
-      val bboxes = BoundingBoxes(fileSchema.getFields.asScala.map(_.getName).flatMap(BoundingBoxField.fromBoundingBox).toSeq)
+      val bboxes = BoundingBoxes(fileSchema.getFields.asScala.map(_.getName).flatMap(BoundingBoxField.fromFieldName).toSeq)
+      val zCols = ZValues(fileSchema.getFields.asScala.map(_.getName).flatMap(ZValueField.fromFieldName).toSeq)
       val visibilities = fileSchema.containsField(VisibilitiesField)
-      val schema = readSchema(fileSchema, sft, readSft, encodings, bboxes, visibilities)
-      SimpleFeatureParquetSchema(readSft.getOrElse(sft), encodings, visibilities, bboxes, metadata.asJava, schema)
+      val schema = readSchema(fileSchema, sft, readSft, encodings, bboxes, zCols, visibilities)
+      SimpleFeatureParquetSchema(readSft.getOrElse(sft), encodings, visibilities, bboxes, zCols, metadata.asJava, schema)
     }
   }
 
@@ -216,43 +223,27 @@ object SimpleFeatureParquetSchema extends LazyLogging {
       readSft: Option[SimpleFeatureType],
       encodings: Encodings,
       bboxes: BoundingBoxes,
+      zValues: ZValues,
       visibilities: Boolean): MessageType = {
-    val consistentSchema = schema(sft, readSft, encodings, bboxes, visibilities) // current schema
-    if (fileSchema.getFieldName(0) == SimpleFeatureParquetSchema.FeatureIdField) {
-      consistentSchema
-    } else {
-      def getMappedField(name: String): Type = consistentSchema.getFields.get(consistentSchema.getFieldIndex(name))
-      // old files - attributes are first, then fid, vis, bboxes
-      val attributes = readSft.getOrElse(sft).getAttributeDescriptors.asScala.map { d =>
-        getMappedField(alphaNumericSafeString(d.getLocalName))
-      }
-      val id = getMappedField(FeatureIdField)
-      val vis = if (!visibilities) { None } else { Some(getMappedField(VisibilitiesField)) }
-      val boxes = sft.getAttributeDescriptors.asScala.flatMap { d =>
-        bboxes.get(d.getLocalName).map(getMappedField)
-      }
-      val fields = attributes ++ Seq(id) ++ vis ++ boxes
-      new MessageType(alphaNumericSafeString(sft.getTypeName), fields.asJava)
-    }
+    schema(sft, encodings, bboxes, zValues, visibilities).toMessageType(readSft.getOrElse(sft), excludeZValues = true)
   }
 
   /**
    * Get the message type for a simple feature type. We need the full sft in order to ensure field ids are
    * consistent, but the schema may be reduced (e.g. on read) based on the filter sft
    *
    * @param sft simple feature type
-   * @param readSft optional feature type used to filter the fields in the schema
    * @param encodings field type encoding
    * @param bboxes include bounding boxes for each row
    * @param visibilities include visibilities
    * @return
    */
   private def schema(
       sft: SimpleFeatureType,
-      readSft: Option[SimpleFeatureType],
       encodings: Encodings,
       bboxes: BoundingBoxes,
-      visibilities: Boolean): MessageType = {
+      zValues: ZValues,
+      visibilities: Boolean): SchemaFields = {
     // note: for iceberg compatibility, field ids need to start at one and increment (without gaps) across all top-level fields.
     // All nested fields (structs, lists) get ids *after* all the top-level fields, once again incrementing without gaps
     val fieldIds = new AtomicInteger(1)
@@ -271,23 +262,18 @@ object SimpleFeatureParquetSchema extends LazyLogging {
     }
     val attributes = {
       val builder = Map.newBuilder[String, Seq[Type]]
-      val nestedFieldIds = new AtomicInteger(fieldIds.get() + sft.getAttributeCount + bboxes.fields.size)
+      val nestedFieldIds = new AtomicInteger(fieldIds.get() + sft.getAttributeCount + bboxes.fields.size + zValues.fields.size)
       sft.getAttributeDescriptors.asScala.foreach { d =>
         val name = alphaNumericSafeString(d.getLocalName)
         val types =
           Seq(buildType(name, ObjectType.selectType(d), encodings, fieldIds, nestedFieldIds)) ++
-            bboxes.get(d.getLocalName).map(bbox => BoundingBoxField.schema(bbox, fieldIds, nestedFieldIds))
+            bboxes.get(d.getLocalName).map(bbox => BoundingBoxField.schema(bbox, fieldIds, nestedFieldIds)) ++
+            zValues.get(d.getLocalName).map(zValue => ZValueField.schema(zValue, fieldIds))
          builder += name -> types
       }
       builder.result()
     }
-    val attributeFields = readSft.getOrElse(sft).getAttributeDescriptors.asScala.flatMap { d =>
-      attributes(alphaNumericSafeString(d.getLocalName))
-    }
-    // note: id field goes at the front of the record, then vis, then attributes and bounding boxes
-    val fields = Seq(id) ++ vis ++ attributeFields
-    val name = alphaNumericSafeString(sft.getTypeName)
-    new MessageType(name, fields.asJava)
+    SchemaFields(id, vis, attributes)
   }
 
   private def icebergSchema(schema: SimpleFeatureParquetSchema): Schema = {
@@ -305,7 +291,8 @@ object SimpleFeatureParquetSchema extends LazyLogging {
       fields.add(NestedField.optional(fieldIds.getAndIncrement(), VisibilitiesField, StringType.get()))
     }
 
-    val nestedFieldIds = new AtomicInteger(fieldIds.get() + schema.sft.getAttributeCount + schema.bboxes.fields.size)
+    val nestedFieldIds =
+      new AtomicInteger(fieldIds.get() + schema.sft.getAttributeCount + schema.bboxes.fields.size + schema.zValues.fields.size)
     schema.sft.getAttributeDescriptors.asScala.foreach { d =>
       val name = alphaNumericSafeString(d.getLocalName)
       aliases.put(name, fieldIds.get())
@@ -314,6 +301,10 @@ object SimpleFeatureParquetSchema extends LazyLogging {
         aliases.put(bbox, fieldIds.get())
         fields.add(BoundingBoxField.icebergSchema(bbox, fieldIds, nestedFieldIds))
       }
+      schema.zValues.get(d.getLocalName).foreach { zValue =>
+        aliases.put(zValue, fieldIds.get())
+        fields.add(ZValueField.icebergSchema(zValue, fieldIds))
+      }
     }
     new Schema(fields, aliases, java.util.Set.of[Integer](1))
   }
@@ -423,4 +414,18 @@ object SimpleFeatureParquetSchema extends LazyLogging {
     }
     typed.build()
   }
+
+  private case class SchemaFields(fid: Type, vis: Option[Type], attributes: Map[String, Seq[Type]]) {
+    def toMessageType(sft: SimpleFeatureType, excludeZValues: Boolean = false): MessageType = {
+      // note: id field goes at the front of the record, then vis, then attributes and bounding boxes
+      val fields = Seq(fid) ++ vis ++ {
+        val all = sft.getAttributeDescriptors.asScala.flatMap { d =>
+          attributes(alphaNumericSafeString(d.getLocalName))
+        }
+        if (!excludeZValues) { all } else { all.filter(f => ZValueField.fromFieldName(f.getName).isEmpty) }
+      }
+      val name = alphaNumericSafeString(sft.getTypeName)
+      new MessageType(name, fields.asJava)
+    }
+  }
 }