Как унифицировать создание DataFrame в databricks-connect и PySpark

Покажем несовместимости создания DataFrame между databricks-connect и PySpark: почему ломается вывод схемы, строгая типизация, как писать стабильные тесты.

Запуск одних и тех же модульных тестов для ETL‑пайплайна через databricks-connect в VS Code и на «чистом» pyspark в Docker-контейнере может выявить тонкие несовместимости. Один из таких случаев — создание DataFrame: код, который без проблем работает с databricks-connect и даже в блокноте Databricks, в pyspark может падать с ошибками вроде CANNOT_INFER_SCHEMA_FOR_TYPE или CANNOT_ACCEPT_OBJECT_IN_TYPE. Источник проблемы — различия в выводе типов и в строгости соблюдения схемы.

Воспроизводим проблему

Сессию Spark инициализируем так, чтобы она работала в обоих окружениях. В VS Code используется databricks-connect; в Docker — локальный сеанс pyspark.

from pyspark.sql import SparkSession

def init_session() -> SparkSession:
    try:
        from databricks.connect import DatabricksSession
        cfg = load_profile()
        return DatabricksSession.builder.remote(
            host=cfg.host,
            token=cfg.token,
            cluster_id=cfg.cluster_id,
        ).getOrCreate()
    except ImportError:
        return (
            SparkSession.builder
            .master("local[4]")
            .config("spark.databricks.clusterUsageTags.clusterAllTags", "{}")
            .getOrCreate()
        )

Первый сбой — создание одноколоночного DataFrame напрямую из списка значений. С databricks-connect это проходит, а в pyspark падает с ошибкой вывода схемы.

spark_ctx = init_session()

frame_values = spark_ctx.createDataFrame(
    [
        "123",
        123456,
        12345678912345,
    ],
    ["my_str_column"],
)

Второй случай — передача целого числа в поле типа FloatType. В databricks-connect это принимается, тогда как pyspark строго отвергает int для FloatType.

from pyspark.sql.types import StructType, StructField, FloatType

rows_raw = [
    (102),
]

schema_spec = StructType([
    StructField("my_float_column", FloatType()),
])

frame_strict = spark_ctx.createDataFrame(rows_raw, schema_spec)

В чём дело на самом деле

Разница сводится к двум моментам. Для одноколоночного ввода databricks-connect принимает «голый» список и выводит схему из одного столбца, а pyspark ожидает, что каждая строка будет кортежем, даже если столбец всего один. Кроме того, при явно заданной схеме databricks-connect терпимее к неявным преобразованиям, тогда как pyspark строго придерживается объявленных типов и требует точного соответствия входных данных.

Решение, работающее в обоих окружениях

Переносимый подход прост: всегда передавайте строки как кортежи, даже если колонка одна, и заранее приводите значения к объявленным типам перед созданием DataFrame.

# Одноколоночный DataFrame: используйте кортежи для каждой строки
stable_one_col = spark_ctx.createDataFrame(
    [
        ("123",),
        (123456,),
        (12345678912345,),
    ],
    ["my_str_column"],
)

stable_one_col.show()

+--------------+
| my_str_column|
+--------------+
|           123|
|        123456|
|12345678912345|
+--------------+

# Соблюдение схемы: приводите к точному ожидаемому типу
from pyspark.sql.types import StructType, StructField, FloatType

rows_casted = [
    (102.0,),
]

schema_spec = StructType([
    StructField("my_float_column", FloatType()),
])

stable_with_schema = spark_ctx.createDataFrame(rows_casted, schema_spec)

stable_with_schema.show()

+---------------+
|my_float_column|
+---------------+
|          102.0|
+---------------+

Почему это важно

Если модульные тесты в GitLab CI запускаются на pyspark, а локальная разработка ведётся с databricks-connect, несогласованное создание DataFrame быстро приводит к нестабильным тестам и ложным падениям. Нормализуйте форму входных данных и явно задавайте типы — это делает тесты детерминированными и позволяет одному и тому же коду без проблем выполняться и в VS Code, и в Docker‑базированном CI. Такой подход также лучше согласуется с тем, как те же синтаксические конструкции исполняются в блокнотах Databricks.

Заключение

При работе сразу с databricks-connect и pyspark используйте согласованный ввод для DataFrame. Для данных с одной колонкой передавайте кортежи, а не голые значения. Если указываете схему, заранее приводите значения к точным целевым типам. Эти два правила убирают наблюдавшиеся расхождения, делают модульные тесты переносимыми между окружениями и уменьшают сюрпризы в CI‑конвейерах.