def test_box2d_top_left(spark: SparkSession):
df = spark.createDataFrame(
[
Row(values=[1.0, 2.0, 1.0, 1.0]),
Row(values=[10.0, 12.0, 1.0, 5.0]),
],
["values"],
).withColumn("bbox", box2d_from_center("values"))
df = df.withColumn("area", area(col("bbox")))
assert_area_equals([1.0, 5.0], df)
def test_box2d_center(spark):
df = spark.createDataFrame(
[
Row(values=[1.5, 2.5, 1.0, 1.0]),
Row(values=[10.5, 14.5, 1.0, 5.0]),
],
["values"],
).withColumn("bbox", box2d_from_center("values"))
df = df.withColumn("area", area(col("bbox")))
assert_area_equals([1.0, 5.0], df)
def test_box2d_udfs(spark):
df = spark.createDataFrame(
[
Row(values=[1.0, 2.0, 2.0, 3.0]),
Row(values=[10.0, 12.0, 11.0, 17.0]),
],
["values"],
).withColumn("bbox", box2d("values"))
df = df.withColumn("area", area(col("bbox")))
assert_area_equals([1.0, 5.0], df)
def test_areas(spark: SparkSession):
"""Test calculating bounding box's area."""
df = spark.createDataFrame(
[
(Box2d(1, 2, 2.0, 3.0), ),
(Box2d(10, 12, 11.0, 17.0), ),
],
["bbox"],
)
df = df.withColumn("area", area(col("bbox")))
assert_area_equals([1.0, 5.0], df)
def test_areas(self):
"""Test calculating bounding box's area."""
df = self.spark.createDataFrame(
[
(Box2d(1, 2, 1.0, 1.0), ),
(Box2d(10, 12, 1.0, 5.0), ),
],
["bbox"],
)
df = df.withColumn("area", area(col("bbox")))
self.assertCountEqual((1.0, 5.0), df.select("area").toPandas()["area"])