Создание схемы данных в Spark

Для оптимизации работы с датафреймом Spark заранее позаботьтесь о задании схемы данных. Это уменьшит время загрузки, сэкономит память и позволит избежать неожиданностей при автоматическом определении типов.

Для примера будем работать с датасетом о пассажирах Титаника. При загрузке данных из источников, не хранящих информацию о типах, по умолчанию колонкам будет присвоен строчный тип:

sdf = spark.read.csv("titanic.csv", header=True)
sdf.printSchema()

Если хотите установить автоматическое определение типов, воспользуйтесь параметром inferSchema:

sdf = spark.read.csv("titanic.csv", header=True, inferSchema=True)
sdf.printSchema()

Типы можно посмотреть и через метод dtypes:

sdf.dtypes

Для задания схемы вручную воспользуйтесь классами типов из модуля pyspark.sql.types:

from pyspark.sql.types import StructType, StructField, StringType,\
                              BooleanType, IntegerType, DoubleType,\
                              ShortType, FloatType

Чаще всего схема определяется как экземпляр объекта StructType, состоящего из списка StructField. Последние требуется инициализировать кортежами с именем поля и его типом для каждой колонки. Также в третьем аргументе конструктора StructField можно задать опцию возможности присутствия пропусков в столбце:

map_l = [('pclass',ShortType()), ('survived',ShortType()), ('name',StringType()),
       ('sex', StringType()), ('age', FloatType()), ('sibsp',ShortType()), 
       ('parch', ShortType()), ('ticket',StringType()), ('fare',FloatType()), 
       ('cabin',StringType()), ('embarked', StringType()), ('boat', StringType()), 
       ('body', FloatType()), ('home.dest', StringType())]
       
sch = StructType([StructField(el[0], el[1], True) for el in map_l])

sdf = spark.read.csv("titanic.csv", header=True, schema=sch)
sdf.dtypes

При создании схемы ориентируйтесь на перечень типов c официальной страницы.