Skip to content

Commit e06bdd5

Browse files
metasimmetasim
committed
Added tests for rf_convert_cell_type and rf_with_no_data.
Fixed a bug with counting data/no_data cells on constant tiles.
1 parent 67108de commit e06bdd5

File tree

6 files changed

+87
-19
lines changed

6 files changed

+87
-19
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ package object expressions {
6868
registry.registerExpression[ExplodeTiles]("rf_explode_tiles")
6969
registry.registerExpression[GetCellType]("rf_cell_type")
7070
registry.registerExpression[SetCellType]("rf_convert_cell_type")
71-
registry.registerExpression[SetNoDataValue]("rf_with_nodata")
71+
registry.registerExpression[SetNoDataValue]("rf_with_no_data")
7272
registry.registerExpression[GetDimensions]("rf_dimensions")
7373
registry.registerExpression[ExtentToGeometry]("st_geometry")
7474
registry.registerExpression[GetGeometry]("rf_geometry")

‎core/src/main/scala/org/locationtech/rasterframes/expressions/tilestats/DataCells.scala‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,12 @@
2121

2222
package org.locationtech.rasterframes.expressions.tilestats
2323

24-
import org.locationtech.rasterframes.expressions.{NullToValue, UnaryRasterOp}
2524
import geotrellis.raster._
26-
import org.apache.spark.sql.{Column, TypedColumn}
27-
import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
2825
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
26+
import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
2927
import org.apache.spark.sql.types.{DataType, LongType}
28+
import org.apache.spark.sql.{Column, TypedColumn}
29+
import org.locationtech.rasterframes.expressions.UnaryRasterOp
3030
import org.locationtech.rasterframes.model.TileContext
3131

3232
@ExpressionDescription(
@@ -40,11 +40,10 @@ import org.locationtech.rasterframes.model.TileContext
4040
357"""
4141
)
4242
case class DataCells(child: Expression) extends UnaryRasterOp
43-
with CodegenFallback with NullToValue {
43+
with CodegenFallback {
4444
override def nodeName: String = "rf_data_cells"
4545
override def dataType: DataType = LongType
4646
override protected def eval(tile: Tile, ctx: Option[TileContext]): Any = DataCells.op(tile)
47-
override def na: Any = 0L
4847
}
4948
object DataCells {
5049
import org.locationtech.rasterframes.encoders.StandardEncoders.PrimitiveEncoders.longEnc

‎core/src/main/scala/org/locationtech/rasterframes/expressions/tilestats/NoDataCells.scala‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,12 @@
2121

2222
package org.locationtech.rasterframes.expressions.tilestats
2323

24-
import org.locationtech.rasterframes.expressions.{NullToValue, UnaryRasterOp}
2524
import geotrellis.raster._
26-
import org.apache.spark.sql.{Column, TypedColumn}
27-
import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
2825
import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback
26+
import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionDescription}
2927
import org.apache.spark.sql.types.{DataType, LongType}
28+
import org.apache.spark.sql.{Column, TypedColumn}
29+
import org.locationtech.rasterframes.expressions.UnaryRasterOp
3030
import org.locationtech.rasterframes.model.TileContext
3131

3232
@ExpressionDescription(
@@ -40,11 +40,10 @@ import org.locationtech.rasterframes.model.TileContext
4040
12"""
4141
)
4242
case class NoDataCells(child: Expression) extends UnaryRasterOp
43-
with CodegenFallback with NullToValue {
43+
with CodegenFallback {
4444
override def nodeName: String = "rf_no_data_cells"
4545
override def dataType: DataType = LongType
4646
override protected def eval(tile: Tile, ctx: Option[TileContext]): Any = NoDataCells.op(tile)
47-
override def na: Any = 0L
4847
}
4948
object NoDataCells {
5049
import org.locationtech.rasterframes.encoders.StandardEncoders.PrimitiveEncoders.longEnc

‎core/src/main/scala/org/locationtech/rasterframes/expressions/transformers/SetNoDataValue.scala‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ import org.locationtech.rasterframes.expressions.{fpTile, row}
4848
)
4949
case class SetNoDataValue(left: Expression, right: Expression) extends BinaryExpression with CodegenFallback with LazyLogging {
5050

51-
override val nodeName: String = "rf_set_nodata"
51+
override val nodeName: String = "rf_with_no_data"
5252
override def dataType: DataType = left.dataType
5353

5454
override def checkInputDataTypes(): TypeCheckResult = {

‎core/src/test/scala/org/locationtech/rasterframes/RasterFunctionsSpec.scala‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,44 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
9696
}
9797
}
9898

99+
describe("cell type operations") {
100+
it("should convert cell type") {
101+
val df = Seq((TestData.injectND(7)(three), TestData.injectND(12)(two))).toDF("three", "two")
102+
103+
val ct = df.select(
104+
rf_convert_cell_type($"three", "uint16ud512") as "three",
105+
rf_convert_cell_type($"two", "float32") as "two"
106+
)
107+
108+
val (ct3, ct2) = ct.as[(Tile, Tile)].first()
109+
110+
ct3.cellType should be (UShortUserDefinedNoDataCellType(512))
111+
ct2.cellType should be (FloatConstantNoDataCellType)
112+
113+
val (cnt3, cnt2) = ct.select(rf_no_data_cells($"three"), rf_no_data_cells($"two")).as[(Long, Long)].first()
114+
115+
cnt3 should be (7)
116+
cnt2 should be (12)
117+
118+
checkDocs("rf_convert_cell_type")
119+
}
120+
it("should change NoData value") {
121+
val df = Seq((TestData.injectND(7)(three), TestData.injectND(12)(two))).toDF("three", "two")
122+
123+
val ct = df.select(
124+
rf_with_no_data($"three", 3) as "three",
125+
rf_with_no_data($"two", 2) as "two"
126+
)
127+
128+
val (cnt3, cnt2) = ct.select(rf_no_data_cells($"three"), rf_no_data_cells($"two")).as[(Long, Long)].first()
129+
130+
cnt3 should be ((cols * rows) - 7)
131+
cnt2 should be ((cols * rows) - 12)
132+
133+
checkDocs("rf_with_no_data")
134+
}
135+
}
136+
99137
describe("arithmetic tile operations") {
100138
it("should local_add") {
101139
val df = Seq((one, two)).toDF("one", "two")
@@ -336,6 +374,25 @@ class RasterFunctionsSpec extends TestEnvironment with RasterMatchers {
336374

337375
checkDocs("rf_no_data_cells")
338376
}
377+
it("should properly count data and nodata cells on constant tiles") {
378+
val rf = Seq(randPRT).toDF("tile")
379+
380+
val df = rf
381+
.withColumn("make", rf_make_constant_tile(99, 3, 4, ByteConstantNoDataCellType))
382+
.withColumn("make2", rf_with_no_data($"make", 99))
383+
384+
df.show(false)
385+
386+
val counts = df.select(
387+
rf_no_data_cells($"make").alias("nodata1"),
388+
rf_data_cells($"make").alias("data1"),
389+
rf_no_data_cells($"make2").alias("nodata2"),
390+
rf_data_cells($"make2").alias("data2")
391+
).as[(Long, Long, Long, Long)].first()
392+
393+
counts should be ((0l, 12l, 12l, 0l))
394+
}
395+
339396
it("should detect no-data tiles") {
340397
val df = Seq(nd).toDF("nd")
341398
df.select(rf_is_no_data_tile($"nd")).first() should be(true)

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

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -253,16 +253,29 @@ def test_exists_for_all(self):
253253

254254
def test_cell_type_in_functions(self):
255255
from pyrasterframes.rf_types import CellType
256-
ct = CellType.float32()
257-
ct_str = CellType.float32().with_no_data_value(-999).cell_type_name # NB int ND into float celltype
256+
ct = CellType.float32().with_no_data_value(-999)
258257

259-
df = self.rf.withColumn('ct_str', rf_convert_cell_type('tile', ct_str)) \
260-
.withColumn('ct', rf_convert_cell_type('tile', ct)) \
261-
.withColumn('make', rf_make_constant_tile(99, 3, 4, CellType.int8()))
258+
df = self.rf.withColumn('ct_str', rf_convert_cell_type('tile', ct.cell_type_name)) \
259+
.withColumn('ct', rf_convert_cell_type('tile', ct)) \
260+
.withColumn('make', rf_make_constant_tile(99, 3, 4, CellType.int8())) \
261+
.withColumn('make2', rf_with_no_data('make', 99))
262262

263-
result = df.select('ct', 'ct_str', 'make').first()
263+
result = df.select('ct', 'ct_str', 'make', 'make2').first()
264264

265265
self.assertEqual(result['ct'].cell_type, ct)
266-
self.assertEqual(result['ct_str'].cell_type, CellType(ct_str))
266+
self.assertEqual(result['ct_str'].cell_type, ct)
267267
self.assertEqual(result['make'].cell_type, CellType.int8())
268+
self.assertEqual(result['make2'].cell_type, CellType.int8().with_no_data_value(99))
269+
270+
counts = df.select(
271+
rf_no_data_cells('make').alias("nodata1"),
272+
rf_data_cells('make').alias("data1"),
273+
rf_no_data_cells('make2').alias("nodata2"),
274+
rf_data_cells('make2').alias("data2")
275+
).first()
276+
277+
self.assertEqual(counts["data1"], 3 * 4)
278+
self.assertEqual(counts["nodata1"], 0)
279+
self.assertEqual(counts["data2"], 0)
280+
self.assertEqual(counts["nodata2"], 3 * 4)
268281

0 commit comments

Comments
 (0)