Merge pull request #57 from s22s/feature/tile-quantile-opts

Alternate approaches to exposing quantile summaries API

Author: vpipkt

core/src/main/scala/org/locationtech/rasterframes/RasterFunctions.scala

@@ -171,9 +171,15 @@ trait RasterFunctions {
   /** Assign a `NoData` value to the tile column. */
   def rf_with_no_data(col: Column, nodata: Column): Column = SetNoDataValue(col, nodata)
 
-  /**  Compute the full column aggregate floating point histogram. */
+  /**  Compute the approximate aggregate floating point histogram using a streaming algorithm, with the default of 80 buckets. */
   def rf_agg_approx_histogram(col: Column): TypedColumn[Any, CellHistogram] = HistogramAggregate(col)
 
+  /**  Compute the approximate aggregate floating point histogram using a streaming algorithm, with the given number of buckets. */
+  def rf_agg_approx_histogram(col: Column, numBuckets: Int): TypedColumn[Any, CellHistogram] = {
+    require(numBuckets > 0, "Must provide a positive number of buckets")
+    HistogramAggregate(col, numBuckets)
+  }
+
   /** Compute the full column aggregate floating point statistics. */
   def rf_agg_stats(col: Column): TypedColumn[Any, CellStatistics] = CellStatsAggregate(col)
 
@@ -186,6 +192,23 @@ trait RasterFunctions {
   /** Computes the number of NoData cells in a column. */
   def rf_agg_no_data_cells(col: Column): TypedColumn[Any, Long] = CellCountAggregate.NoDataCells(col)
 
+  /**
+    * Calculates the approximate quantiles of a tile column of a DataFrame.
+    * @param tile tile column to extract cells from.
+    * @param probabilities a list of quantile probabilities
+    *   Each number must belong to [0, 1].
+    *   For example 0 is the minimum, 0.5 is the median, 1 is the maximum.
+    * @param relativeError The relative target precision to achieve (greater than or equal to 0).
+    * @return the approximate quantiles at the given probabilities of each column
+    */
+  def rf_agg_approx_quantiles(
+    tile: Column,
+    probabilities: Seq[Double],
+    relativeError: Double = 0.00001): TypedColumn[Any, Seq[Double]] = {
+    require(probabilities.nonEmpty, "at least one quantile probability is required")
+    ApproxCellQuantilesAggregate(tile, probabilities, relativeError)
+  }
+
   /** Compute the Tile-wise mean */
   def rf_tile_mean(col: Column): TypedColumn[Any, Double] =
     TileMean(col)
@@ -546,14 +569,17 @@ trait RasterFunctions {
   /** Return the incoming tile untouched. */
   def rf_identity(tileCol: Column): Column = Identity(tileCol)
 
-  /** Create a row for each cell in Tile. */
+  /** Create a row for each cell in Tile.
+    * The output will include the columns `column_index`, `row_index` indicating where in the tile the cell originated. */
   def rf_explode_tiles(cols: Column*): Column = rf_explode_tiles_sample(1.0, None, cols: _*)
 
-  /** Create a row for each cell in Tile with random sampling and optional seed. */
+  /** Create a row for each cell in Tile with random sampling and optional seed.
+    * The output will include the columns `column_index`, `row_index` indicating where in the tile the cell originated. */
   def rf_explode_tiles_sample(sampleFraction: Double, seed: Option[Long], cols: Column*): Column =
     ExplodeTiles(sampleFraction, seed, cols)
 
-  /** Create a row for each cell in Tile with random sampling (no seed). */
+  /** Create a row for each cell in Tile with random sampling (no seed).
+    * The output will include the columns `column_index`, `row_index` indicating where in the tile the cell originated. */
   def rf_explode_tiles_sample(sampleFraction: Double, cols: Column*): Column =
     ExplodeTiles(sampleFraction, None, cols)
 }

core/src/main/scala/org/locationtech/rasterframes/encoders/StandardSerializers.scala

@@ -21,17 +21,21 @@
 
 package org.locationtech.rasterframes.encoders
 
+import java.nio.ByteBuffer
+
 import com.github.blemale.scaffeine.Scaffeine
 import geotrellis.proj4.CRS
 import geotrellis.raster._
 import geotrellis.spark._
 import geotrellis.spark.tiling.LayoutDefinition
 import geotrellis.vector._
+import org.apache.spark.sql.catalyst.util.QuantileSummaries
 import org.apache.spark.sql.types._
 import org.locationtech.jts.geom.Envelope
 import org.locationtech.rasterframes.TileType
 import org.locationtech.rasterframes.encoders.CatalystSerializer.{CatalystIO, _}
 import org.locationtech.rasterframes.model.LazyCRS
+import org.locationtech.rasterframes.util.KryoSupport
 
 /** Collection of CatalystSerializers for third-party types. */
 trait StandardSerializers {
@@ -294,9 +298,23 @@ trait StandardSerializers {
   implicit val spatialKeyTLMSerializer = tileLayerMetadataSerializer[SpatialKey]
   implicit val spaceTimeKeyTLMSerializer = tileLayerMetadataSerializer[SpaceTimeKey]
 
+  implicit val quantileSerializer: CatalystSerializer[QuantileSummaries] = new CatalystSerializer[QuantileSummaries] {
+    override val schema: StructType = StructType(Seq(
+      StructField("quantile_serializer_kryo", BinaryType, false)
+    ))
+
+    override protected def to[R](t: QuantileSummaries, io: CatalystSerializer.CatalystIO[R]): R = {
+      val buf = KryoSupport.serialize(t)
+      io.create(buf.array())
+    }
+
+    override protected def from[R](t: R, io: CatalystSerializer.CatalystIO[R]): QuantileSummaries = {
+      KryoSupport.deserialize[QuantileSummaries](ByteBuffer.wrap(io.getByteArray(t, 0)))
+    }
+  }
 }
 
-object StandardSerializers {
+object StandardSerializers extends StandardSerializers {
   private val s2ctCache = Scaffeine().build[String, CellType](
     (s: String) => CellType.fromName(s)
   )

core/src/main/scala/org/locationtech/rasterframes/expressions/aggregates/ApproxCellQuantilesAggregate.scala

@@ -0,0 +1,88 @@
+/*
+ * This software is licensed under the Apache 2 license, quoted below.
+ *
+ * Copyright 2019 Astraea, Inc.
+ *
+ * Licensed 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package org.locationtech.rasterframes.expressions.aggregates
+
+import geotrellis.raster.{Tile, isNoData}
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.catalyst.util.QuantileSummaries
+import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
+import org.apache.spark.sql.{Column, Encoder, Row, TypedColumn, types}
+import org.apache.spark.sql.types.{DataTypes, StructField, StructType}
+import org.locationtech.rasterframes.TileType
+import org.locationtech.rasterframes.encoders.CatalystSerializer._
+import org.locationtech.rasterframes.expressions.accessors.ExtractTile
+
+
+case class ApproxCellQuantilesAggregate(probabilities: Seq[Double], relativeError: Double) extends UserDefinedAggregateFunction {
+  import org.locationtech.rasterframes.encoders.StandardSerializers.quantileSerializer
+
+  override def inputSchema: StructType = StructType(Seq(
+    StructField("value", TileType, true)
+  ))
+
+  override def bufferSchema: StructType = StructType(Seq(
+    StructField("buffer", schemaOf[QuantileSummaries], false)
+  ))
+
+  override def dataType: types.DataType = DataTypes.createArrayType(DataTypes.DoubleType)
+
+  override def deterministic: Boolean = true
+
+  override def initialize(buffer: MutableAggregationBuffer): Unit =
+    buffer.update(0, new QuantileSummaries(QuantileSummaries.defaultCompressThreshold, relativeError).toRow)
+
+  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
+    val qs = buffer.getStruct(0).to[QuantileSummaries]
+    if (!input.isNullAt(0)) {
+      val tile = input.getAs[Tile](0)
+      var result = qs
+      tile.foreachDouble(d => if (!isNoData(d)) result = result.insert(d))
+      buffer.update(0, result.toRow)
+    }
+    else buffer
+  }
+
+  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
+    val left = buffer1.getStruct(0).to[QuantileSummaries]
+    val right = buffer2.getStruct(0).to[QuantileSummaries]
+    val merged = left.compress().merge(right.compress())
+    buffer1.update(0, merged.toRow)
+  }
+
+  override def evaluate(buffer: Row): Seq[Double] = {
+    val summaries = buffer.getStruct(0).to[QuantileSummaries]
+    probabilities.flatMap(summaries.query)
+  }
+}
+
+object ApproxCellQuantilesAggregate {
+  private implicit def doubleSeqEncoder: Encoder[Seq[Double]] = ExpressionEncoder()
+
+  def apply(
+    tile: Column,
+    probabilities: Seq[Double],
+    relativeError: Double = 0.00001): TypedColumn[Any, Seq[Double]] = {
+    new ApproxCellQuantilesAggregate(probabilities, relativeError)(ExtractTile(tile))
+      .as(s"rf_agg_approx_quantiles")
+      .as[Seq[Double]]
+  }
+}

core/src/main/scala/org/locationtech/rasterframes/expressions/aggregates/HistogramAggregate.scala

@@ -98,7 +98,10 @@ object HistogramAggregate {
   import org.locationtech.rasterframes.encoders.StandardEncoders.cellHistEncoder
 
   def apply(col: Column): TypedColumn[Any, CellHistogram] =
-    new HistogramAggregate()(ExtractTile(col))
+    apply(col, StreamingHistogram.DEFAULT_NUM_BUCKETS)
+
+  def apply(col: Column, numBuckets: Int): TypedColumn[Any, CellHistogram] =
+    new HistogramAggregate(numBuckets)(ExtractTile(col))
       .as(s"rf_agg_approx_histogram($col)")
       .as[CellHistogram]
 

core/src/main/scala/org/locationtech/rasterframes/extensions/DataFrameMethods.scala

@@ -155,9 +155,6 @@ trait DataFrameMethods[DF <: DataFrame] extends MethodExtensions[DF] with Metada
   def withPrefixedColumnNames(prefix: String): DF =
     self.columns.foldLeft(self)((df, c) ⇒ df.withColumnRenamed(c, s"$prefix$c").asInstanceOf[DF])
 
-  /**  */
-  def tileStat(): RasterFrameStatFunctions = new RasterFrameStatFunctions(self)
-
   /**
     * Performs a jeft join on the dataframe `right` to this one, reprojecting and merging tiles as necessary.
     * The operation is logically a "left outer" join, with the left side also determining the target CRS and extents.