Apache Spark давно перестал быть инструментом исключительно для кластеров из десятков машин — его прекрасно можно запустить локально, прямо в блокноте Google Colab, и пройти на нём весь путь от сырых данных до обученной модели. В этом разборе показан именно такой сценарий: как с помощью PySpark построить сквозной пайплайн — от создания DataFrame и SQL-запросов до оконных функций, join'ов, обучения модели логистической регрессии и сохранения результатов в формате Parquet.
Подготовка окружения и первый DataFrame
Работа начинается с установки PySpark и запуска локальной Spark-сессии с ограничением числа shuffle-партиций до четырёх — это ускоряет операции на небольшом наборе данных. Дальше создаётся структурированный DataFrame с информацией о восьми пользователях: страна, дата регистрации, доход и тип подписки (basic, standard, premium).
!pip install -q pyspark==3.5.1
from pyspark.sql import SparkSession, functions as F, Window
from pyspark.sql.types import IntegerType, StringType, StructType, StructField, FloatType
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.evaluation import MulticlassClassificationEvaluator
spark = (SparkSession.builder.appName("ColabSparkAdvancedTutorial")
.master("local[*]")
.config("spark.sql.shuffle.partitions", "4")
.getOrCreate())
print("Spark version:", spark.version)
data = [
(1, "Alice", "IN", "2025-10-01", 56000.0, "premium"),
(2, "Bob", "US", "2025-10-03", 43000.0, "standard"),
(3, "Carlos", "IN", "2025-09-27", 72000.0, "premium"),
(4, "Diana", "UK", "2025-09-30", 39000.0, "standard"),
(5, "Esha", "IN", "2025-10-02", 85000.0, "premium"),
(6, "Farid", "AE", "2025-10-02", 31000.0, "basic"),
(7, "Gita", "IN", "2025-09-29", 46000.0, "standard"),
(8, "Hassan", "PK", "2025-10-01", 52000.0, "premium"),
]
schema = StructType([
StructField("id", IntegerType(), False),
StructField("name", StringType(), True),
StructField("country", StringType(), True),
StructField("signup_date", StringType(), True),
StructField("income", FloatType(), True),
StructField("plan", StringType(), True),
])
df = spark.createDataFrame(data, schema)
df.show()
Этот блок закладывает основу для всех последующих преобразований: структурированные данные о пользователях становятся отправной точкой для аналитики и обучения модели.
Трансформации, SQL-запросы и оконные функции
Дальше в датасет добавляются новые колонки — временная метка регистрации, год, месяц и флаг «пользователь из Индии». Полученная таблица регистрируется как временное SQL-представление, по которому выполняется агрегирующий запрос: количество пользователей и средний доход по странам. Затем с помощью оконной функции rank() пользователи внутри каждой страны ранжируются по доходу, а пользовательская функция (UDF) присваивает каждому тарифному плану числовой приоритет.
df2 = (df.withColumn("signup_ts", F.to_timestamp("signup_date"))
.withColumn("year", F.year("signup_ts"))
.withColumn("month", F.month("signup_ts"))
.withColumn("is_india", (F.col("country") == "IN").cast("int")))
df2.show()
df2.createOrReplaceTempView("users")
spark.sql("""
SELECT country, COUNT(*) AS cnt, AVG(income) AS avg_income
FROM users
GROUP BY country
ORDER BY cnt DESC
""").show()
w = Window.partitionBy("country").orderBy(F.col("income").desc())
df_ranked = df2.withColumn("income_rank_in_country", F.rank().over(w))
df_ranked.show()
def plan_priority(plan):
if plan == "premium": return 3
if plan == "standard": return 2
if plan == "basic": return 1
return 0
plan_priority_udf = F.udf(plan_priority, IntegerType())
df_udf = df_ranked.withColumn("plan_priority", plan_priority_udf(F.col("plan")))
df_udf.show()
Этот этап показывает, что Spark умеет одинаково удобно работать и как SQL-движок, и как API для DataFrame — переключаться между ними можно без потери производительности.
Объединение данных и агрегация по регионам
Чтобы обогатить пользовательские данные контекстом, к ним присоединяется вторая таблица — справочник стран с указанием региона и численности населения (в миллиардах). После join'а данные группируются по региону и тарифному плану, для каждой комбинации считаются количество пользователей и средний доход.
country_data = [
("IN", "Asia", 1.42), ("US", "North America", 0.33),
("UK", "Europe", 0.07), ("AE", "Asia", 0.01), ("PK", "Asia", 0.24),
]
country_schema = StructType([
StructField("country", StringType(), True),
StructField("region", StringType(), True),
StructField("population_bn", FloatType(), True),
])
country_df = spark.createDataFrame(country_data, country_schema)
joined = df_udf.alias("u").join(country_df.alias("c"), on="country", how="left")
joined.show()
region_stats = (joined.groupBy("region", "plan")
.agg(F.count("*").alias("users"),
F.round(F.avg("income"), 2).alias("avg_income"))
.orderBy("region", "plan"))
region_stats.show()
Это классический паттерн обогащения данных справочниками — Spark справляется с ним так же легко, как с одной таблицей, даже если справочник значительно больше по объёму.
Машинное обучение: логистическая регрессия
Обогащённый датасет становится обучающей выборкой для простой задачи бинарной классификации — предсказать, оформит ли пользователь премиум-подписку. Категориальный признак «страна» индексируется через StringIndexer, а числовые признаки (доход, индекс страны, приоритет плана) собираются в вектор features с помощью VectorAssembler. Данные делятся на обучающую и тестовую выборки в пропорции 70/30, после чего на них обучается модель логистической регрессии из Spark MLlib.
ml_df = joined.withColumn("label", (F.col("plan") == "premium").cast("int")).na.drop()
country_indexer = StringIndexer(inputCol="country", outputCol="country_idx", handleInvalid="keep")
country_fitted = country_indexer.fit(ml_df)
ml_df2 = country_fitted.transform(ml_df)
assembler = VectorAssembler(inputCols=["income", "country_idx", "plan_priority"], outputCol="features")
ml_final = assembler.transform(ml_df2)
train_df, test_df = ml_final.randomSplit([0.7, 0.3], seed=42)
lr = LogisticRegression(featuresCol="features", labelCol="label", maxIter=20)
lr_model = lr.fit(train_df)
preds = lr_model.transform(test_df)
preds.select("name", "country", "income", "plan", "label", "prediction", "probability").show(truncate=False)
evaluator = MulticlassClassificationEvaluator(labelCol="label", predictionCol="prediction", metricName="accuracy")
acc = evaluator.evaluate(preds)
print("Classification accuracy:", acc)
Точность модели выводится через MulticlassClassificationEvaluator — так наглядно видно, что MLlib встраивается в общий пайплайн без дополнительных инструментов и переключений между фреймворками.
Сохранение в Parquet и анализ плана запроса
Финальный шаг — запись обработанных данных в колоночный формат Parquet и их обратное чтение для проверки целостности. Дополнительно выполняется SQL-запрос по недавним регистрациям (начиная с 1 октября 2025 года), а метод explain() показывает план выполнения запроса — что полезно для оптимизации на реальных объёмах данных. В конце Spark-сессия корректно останавливается вызовом spark.stop().
output_path = "/content/spark_users_parquet"
joined.write.mode("overwrite").parquet(output_path)
parquet_df = spark.read.parquet(output_path)
print("Parquet reloaded:")
parquet_df.show()
recent = spark.sql("""
SELECT name, country, income, signup_ts
FROM users
WHERE signup_ts >= '2025-10-01'
ORDER BY signup_ts DESC
""")
recent.show()
recent.explain()
spark.stop()
Что это даёт на практике
Весь пример укладывается в один блокнот Colab, но демонстрирует полный набор инструментов Spark: DataFrame API, Spark SQL, оконные функции, UDF, join и агрегации, MLlib для обучения модели и Parquet для персистентности данных. Для команд, которые хотят прототипировать data-пайплайны перед тем, как разворачивать их на распределённом кластере, это удобный способ отработать логику локально — код почти без изменений переносится на production-инсталляцию Spark.
Источник: MarkTechPost.