Add documentation, python bindings for rf_agg_approx_quantiles

Attempt to register with SQL

Signed-off-by: Jason T. Brown <jason@astraea.earth>

Author: vpipkt

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

@@ -23,7 +23,10 @@ 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.expressions.aggregate.{AggregateExpression, AggregateFunction, AggregateMode, Complete}
+import org.apache.spark.sql.catalyst.expressions.{ExprId, Expression, ExpressionDescription, NamedExpression}
 import org.apache.spark.sql.catalyst.util.QuantileSummaries
+import org.apache.spark.sql.execution.aggregate.ScalaUDAF
 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}
@@ -85,4 +88,29 @@ object ApproxCellQuantilesAggregate {
       .as(s"rf_agg_approx_quantiles")
       .as[Seq[Double]]
   }
+
+  /** Adapter hack to allow UserDefinedAggregateFunction to be referenced as an expression. */
+  @ExpressionDescription(
+    usage = "_FUNC_(tile, probabilities, relativeError) - Compute aggregate cell histogram over a tile column.",
+    arguments = """
+  Arguments:
+    * tile - tile column to analyze
+    * probabilities - array of double values in [0, 1] at which to compute quantiles
+    * relativeError - non-negative error tolerance""",
+    examples = """
+  Examples:
+    > SELECT _FUNC_(tile, array(0.1, 0.25, 0.5, 0.75, 0.9), 0.001);
+      ..."""
+  )
+  class ApproxCellQuantilesUDAF(aggregateFunction: AggregateFunction, mode: AggregateMode, isDistinct: Boolean, resultId: ExprId)
+    extends AggregateExpression(aggregateFunction, mode, isDistinct, resultId) {
+    def this(child: Expression, probabilities: Seq[Double], relativeError: Double) =
+      this(ScalaUDAF(Seq(ExtractTile(child)), new ApproxCellQuantilesAggregate(probabilities, relativeError)), Complete, false, NamedExpression.newExprId)
+    override def nodeName: String = "rf_agg_approx_quantiles"
+  }
+
+  object ApproxCellQuantilesUDAF {
+    def apply(child: Expression, probabilities: Seq[Double], relativeError: Double): ApproxCellQuantilesUDAF = new ApproxCellQuantilesUDAF(child, probabilities, relativeError)
+    def apply(child: Expression, probabilities: Seq[Double]): ApproxCellQuantilesUDAF = new ApproxCellQuantilesUDAF(child, probabilities, 0.00001)
+  }
 }

core/src/main/scala/org/locationtech/rasterframes/expressions/package.scala

@@ -117,6 +117,7 @@ package object expressions {
     registry.registerExpression[CellCountAggregate.NoDataCells]("rf_agg_no_data_cells")
     registry.registerExpression[CellStatsAggregate.CellStatsAggregateUDAF]("rf_agg_stats")
     registry.registerExpression[HistogramAggregate.HistogramAggregateUDAF]("rf_agg_approx_histogram")
+    registry.registerExpression[ApproxCellQuantilesAggregate.ApproxCellQuantilesUDAF]("rf_agg_approx_quantiles")
     registry.registerExpression[LocalStatsAggregate.LocalStatsAggregateUDAF]("rf_agg_local_stats")
     registry.registerExpression[LocalTileOpAggregate.LocalMinUDAF]("rf_agg_local_min")
     registry.registerExpression[LocalTileOpAggregate.LocalMaxUDAF]("rf_agg_local_max")

core/src/test/scala/org/locationtech/rasterframes/RasterFramesStatsSpec.scala

@@ -71,6 +71,17 @@ class RasterFramesStatsSpec extends TestEnvironment with TestData {
       // computing externally with numpy we arrive at 7963, 10068, 12160 for these quantiles
       result should contain inOrderOnly(7963.0, 10068.0, 12160.0)
     }
+
+    it("should compute approx percentiles with SQL") {
+      val result = df.selectExpr("rf_agg_approx_quantiles(tile, array(0.1, 0.5, 0.9), 0.00001) as iles")
+        .first()
+        .getSeq[Double](0)
+
+      result.length should be (3)
+
+      // computing externally with numpy we arrive at 7963, 10068, 12160 for these quantiles
+      result should contain inOrderOnly(7963.0, 10068.0, 12160.0)
+    }
   }
 }
 

docs/src/main/paradox/reference.md

@@ -628,6 +628,11 @@ Aggregates over the `tile` and returns statistical summaries of cell values: num
 
 Aggregates over all of the rows in DataFrame of `tile` and returns a count of each cell value to create a histogram with values are plotted on the x-axis and counts on the y-axis. Related is the @ref:[`rf_tile_histogram`](reference.md#rf-tile-histogram) function which operates on a single row at a time.
 
+### rf_agg_approx_quantiles
+
+    Array[Double] rf_agg_approx_quantiles(Tile tile, List[float] probabilities, float relative_error)
+    
+Calculates the approximate quantiles of a tile column of a DataFrame. `probabilities` is a list of float values at which to compute the quantiles. These must belong to [0, 1]. For example 0 is the minimum, 0.5 is the median, 1 is the maximum. Returns an array of values approximately at the specified `probabilities`.
 
 ## Tile Local Aggregate Statistics
 

docs/src/main/paradox/release-notes.md

@@ -9,6 +9,7 @@
 * Added `withSpatialIndex` to RasterSourceDataSource to pre-partition tiles based on tile extents mapped to a Z2 space-filling curve 
 * Add `rf_mask_by_bit`, `rf_mask_by_bits` and `rf_local_extract_bits` to deal with bit packed quality masks. Updated the masking documentation to demonstrate the use of these functions.
 * Throw an `IllegalArgumentException` when attempting to apply a mask to a `Tile` whose `CellType` has no NoData defined. ([#409](https://github.com/locationtech/rasterframes/issues/384))
+* Add `rf_agg_approx-quantiles` function to compute cell quantiles across an entire column.
 
 ### 0.8.4
 

pyrasterframes/src/main/python/pyrasterframes/rasterfunctions.py

@@ -313,6 +313,22 @@ def rf_agg_approx_histogram(tile_col):
     return _apply_column_function('rf_agg_approx_histogram', tile_col)
 
 
+def rf_agg_approx_quantiles(tile_col, probabilities, relative_error=0.00001):
+    """
+    Calculates the approximate quantiles of a tile column of a DataFrame.
+
+    :param tile_col: 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 relative_error: The relative target precision to achieve (greater than or equal to 0). Default is 0.00001
+    :return: An array of values approximately at the specified `probabilities`
+    """
+
+    _jfn = RFContext.active().lookup('rf_agg_approx_quantiles')
+    _tile_col = _to_java_column(tile_col)
+    return Column(_jfn(_tile_col, probabilities, relative_error))
+
+
 def rf_agg_stats(tile_col):
     """Compute the full column aggregate floating point statistics"""
     return _apply_column_function('rf_agg_stats', tile_col)

pyrasterframes/src/main/python/tests/RasterFunctionsTests.py

@@ -25,7 +25,7 @@
 from pyspark.sql.functions import *
 
 import numpy as np
-from numpy.testing import assert_equal
+from numpy.testing import assert_equal, assert_allclose
 
 from unittest import skip
 from . import TestEnvironment
@@ -133,6 +133,12 @@ def test_aggregations(self):
         self.assertEqual(row['rf_agg_no_data_cells(tile)'], 1000)
         self.assertEqual(row['rf_agg_stats(tile)'].data_cells, row['rf_agg_data_cells(tile)'])
 
+    def test_agg_approx_quantiles(self):
+        agg = self.rf.agg(rf_agg_approx_quantiles('tile', [0.1, 0.5, 0.9, 0.98]))
+        result = agg.first()[0]
+        # expected result from computing in external python process
+        assert_allclose(result, np.array([7412., 7638., 7671., 7675.]))
+
     def test_sql(self):
 
         self.rf.createOrReplaceTempView("rf_test_sql")