使用Databricks的单元测试第一部分 - 使用Databricks连接的PySpark单元测试

方石剑
293 阅读6分钟

用Databricks进行单元测试

第一部分 - 使用Databricks Connect进行PySpark单元测试

在我最近的项目中，我第一次使用Databricks工作。

起初，我发现使用Databricks来编写生产代码有些刺耳--使用网络门户中的笔记本并不是对开发者最友好的，我发现它类似于使用Jupyter笔记本来编写生产代码。

然而，改变游戏规则的是：进入Databricks Connect，一种在Databricks集群上远程执行代码的方式。我回到了家里，在舒适的IDE中开发，在云中运行PySpark命令。我现在真的很喜欢使用Databricks，并且很乐意向任何需要在大数据上进行分布式数据工程的人推荐它。

如果你想在Databricks中自己编写生产代码，你肯定需要做的一件事就是单元测试。这篇博文以及下一部分，旨在通过PySpark中单元测试功能的一个超级简单的例子来帮助你做到这一点。

要跟上这篇博文，你需要

  • Python 3.7
  • 微软Azure中的Databricks工作区，以及一个运行Databricks Runtime 7.3 LTS的集群

快速免责声明：在写作时，我目前是一名微软员工。

设置你的本地环境

首先，在这里克隆与这篇博文相配套的资源库。

现在创建一个新的虚拟环境并运行。

pip install -r requirements.txt

这将安装一些测试要求、Databricks Connect和我们git仓库中定义的python包。

然后你要设置你的Databricks Connect。你可以通过运行databricks-connect configure ，并按照Databricks Connect文档中给出的指示来完成。

你可以通过运行来测试你的Databricks Connect是否正常工作。

databricks-connect test

测试的功能

我们将测试一个函数，它以Spark DataFrame的形式接收一些数据，并以Spark DataFrame的形式返回一些转换后的数据。

假设我们从一些看起来像这样的数据开始，我们有3个泵在泵送液体。

pump_idstart_time结束时间泵送的升数
12021-02-01 01:05:322021-02-01 01:09:1324
22021-02-01 01:09:142021-02-01 01:14:1741
12021-02-01 01:14:182021-02-01 01:15:5811
22021-02-01 01:15:592021-02-01 01:18:2613
12021-02-01 01:18:272021-02-01 01:26:2645
32021-02-01 01:26:272021-02-01 01:38:5715

而我们想知道这些泵中的每一个平均每秒抽出的升数。("litres "的英式英语拼写😉)。

所以我们想要一个类似这样的输出。

pump_id总时间_秒泵送的总升数每秒平均升数
1800800.12
2450540.2
3750150.02

我们可以创建一个如下的函数来完成这个任务。

from pyspark.sql.functions import col, sum as col_sum


def aggregate_pump_data(pump_data_df):
    aggregated_df = pump_data_df.withColumn(
        "duration_seconds",
        (
            col("end_time").cast('timestamp').cast('long')
            - col("start_time").cast('timestamp').cast('long')
        )
    ).groupBy("pump_id").agg(
        col_sum("duration_seconds").alias("total_duration_seconds"),
        col_sum('litres_pumped').alias("total_litres_pumped")
    ).withColumn(
        "avg_litres_per_second",
        col("total_litres_pumped") / col("total_duration_seconds")
    )
    return aggregated_df

这个函数可以在我们的资源库中找到databricks_pkg/databricks_pkg/pump_utils.py 。

单元测试我们的函数

我们函数的单元测试可以在资源库中找到：databricks_pkg/test/test_pump_utils.py 。

要运行单元测试，请运行。

pytest -v ./databricks_pkg/test/test_pump_utils.py

如果一切工作正常，单元测试应该通过。

首先，我将展示单元测试文件，然后逐行浏览它。

import unittest

import pandas as pd
from pyspark.sql import Row, SparkSession
from pyspark.sql.dataframe import DataFrame

from databricks_pkg.pump_utils import get_litres_per_second


class TestGetLitresPerSecond(unittest.TestCase):
    def setUp(self):
        self.spark = SparkSession.builder.getOrCreate()

    def test_get_litres_per_second(self):
        test_data = [
            # pump_id, start_time, end_time, litres_pumped
            (1, '2021-02-01 01:05:32', '2021-02-01 01:09:13', 24),
            (2, '2021-02-01 01:09:14', '2021-02-01 01:14:17', 41),
            (1, '2021-02-01 01:14:18', '2021-02-01 01:15:58', 11),
            (2, '2021-02-01 01:15:59', '2021-02-01 01:18:26', 13),
            (1, '2021-02-01 01:18:27', '2021-02-01 01:26:26', 45),
            (3, '2021-02-01 01:26:27', '2021-02-01 01:38:57', 15)
        ]
        test_data = [
            {
                'pump_id': row[0],
                'start_time': row[1],
                'end_time': row[2],
                'litres_pumped': row[3]
            } for row in test_data
        ]
        test_df = self.spark.createDataFrame(map(lambda x: Row(**x), test_data))
        output_df = get_litres_per_second(test_df)

        self.assertIsInstance(output_df, DataFrame)

        output_df_as_pd = output_df.sort('pump_id').toPandas()

        expected_output_df = pd.DataFrame([
            {
                'pump_id': 1,
                'total_duration_seconds': 800,
                'total_litres_pumped': 80,
                'avg_litres_per_second': 0.1
            },
            {
                'pump_id': 2,
                'total_duration_seconds': 450,
                'total_litres_pumped': 54,
                'avg_litres_per_second': 0.12
            },
            {
                'pump_id': 3,
                'total_duration_seconds': 750,
                'total_litres_pumped': 15,
                'avg_litres_per_second': 0.02
            },
        ])
        pd.testing.assert_frame_equal(expected_output_df, output_df_as_pd)

现在我们将逐行浏览这个文件。

输入

单元测试功能从一些进口开始，我们从内置包开始，然后是外部包，最后是内部包--其中包括我们要测试的功能。

import unittest

import pandas as pd
from pyspark.sql import Row, SparkSession
from pyspark.sql.dataframe import DataFrame

from databricks_pkg.pump_utils import get_litres_per_second

定义测试类

测试类是unittest.TestCase 类的一个子类。unittest.TestCase 自带了一系列的类方法，用于帮助单元测试。

虽然在本例中我们只运行一个测试，但我们可能已经运行了多个测试，我们可以使用相同的TestGetLitresPerSecond.spark 属性作为我们的火花会话。

class TestGetLitresPerSecond(unittest.TestCase):
    def setUp(self):
        self.spark = SparkSession.builder.getOrCreate()

定义测试数据

接下来我们要做的是定义一些测试数据，我们将使用本文章顶部显示的测试数据。它首先被定义为一个图元的列表，然后我用一个列表理解法将其转换为一个dicts的列表。

我这样定义是为了方便阅读，你可以随心所欲地定义你的测试数据。

def test_get_litres_per_second(self):
    test_data = [
        # pump_id, start_time, end_time, litres_pumped
        (1, '2021-02-01 01:05:32', '2021-02-01 01:09:13', 24),
        (2, '2021-02-01 01:09:14', '2021-02-01 01:14:17', 41),
        (1, '2021-02-01 01:14:18', '2021-02-01 01:15:58', 11),
        (2, '2021-02-01 01:15:59', '2021-02-01 01:18:26', 13),
        (1, '2021-02-01 01:18:27', '2021-02-01 01:26:26', 45),
        (3, '2021-02-01 01:26:27', '2021-02-01 01:38:57', 15)
    ]
    test_data = [
        {
            'pump_id': row[0],
            'start_time': row[1],
            'end_time': row[2],
            'litres_pumped': row[3]
        } for row in test_data
   ]

转换为Spark DataFrame

这里我们首先要使用我们的spark会话在我们的Databricks集群中运行，这将我们的dict列表转换为spark DataFrame。

test_df = self.spark.createDataFrame(map(lambda x: Row(**x), test_data))

运行我们的函数

现在我们用我们的测试DataFrame运行我们要测试的函数。

output_df = get_litres_per_second(test_df)

测试该函数的输出

首先要检查的是我们函数的输出是否是我们期望的正确数据类型，我们可以使用unittest.TestCase 类的方法assertIsInstance 。

self.assertIsInstance(output_df, DataFrame)

然后我们将把我们的spark DataFrame转换成pandas DataFrame。我们还需要对DataFrame进行排序，不能保证DataFrame的处理输出是按任何顺序进行的，尤其是行被分割并在不同的节点上处理。

output_df_as_pd = output_df.sort('pump_id').toPandas()

然后我们可以检查这个输出的DataFrame是否等于我们的预期输出。

expected_output_df = pd.DataFrame([
    {
        'pump_id': 1,
        'total_duration_seconds': 800,
        'total_litres_pumped': 80,
        'avg_litres_per_second': 0.1
    },
    {
        'pump_id': 2,
        'total_duration_seconds': 450,
        'total_litres_pumped': 54,
        'avg_litres_per_second': 0.12
    },
    {
        'pump_id': 3,
        'total_duration_seconds': 750,
        'total_litres_pumped': 15,
        'avg_litres_per_second': 0.02
    },
])
pd.testing.assert_frame_equal(expected_output_df, output_df_as_pd)

希望这篇博文能帮助你了解使用Databricks和Databricks Connect进行PySpark单元测试的基本知识。

在本系列博文的下一部分，我们将深入探讨如何将这个单元测试整合到我们的CI管道中。

第二部分请看这里。

avatar
文章
阅读
粉丝