pysparkチートシート

# spdf: spark dataframe
columns = spdf[99:999]
columns_fixed = ['`' + column + '`' for column in columns ]
spdf_tmp = spdf.select(columns_fixed)

def printDf(sprkDF): 
    # format Spark Dataframe like pandas dataframe
    # sparkのdataframeをpandasのdataframeのように整形して出力する
    newdf = sprkDF.toPandas()
    from IPython.core.display import display, HTML
    return HTML(newdf.to_html())

df2 = df.select(*(F.col(c).cast("long").alias(c) for c in df.columns))

import pyspark.sql.functions as F
from pyspark.sql import Window

# spdf: spark dataframe
spdf = spdf.withColumn('new_column',F.row_number().over(Window.partitionBy('partition_column').orderBy('order_by_column')))

import pyspark.sql.functions as F

spdf = spdf.withColumn('time_column', F.to_timestamp(F.col('time_column')))

import pyspark.sql.functions as F

start = '2020-11-01 00:00:00'
end = '2021-02-01 00:00:00'

# 複数条件ある時に時々変な不具合防ぐために条件ごとに括弧で括るの推奨
spdf = spdf.filter( (F.col('Createdat') >= start) & (F.col('Createdat') <= end) )

column_name = 'Createdat'
spdf.select(column_name).agg(\
    F.min(F.col(column_name))\
    , F.max(F.col(column_name))\
    , F.avg(F.col(column_name))\
    , F.sum(F.col(column_name))\
    ).show()

import pyspark.sql.functions as F
import pyspark.sql.types as T
from pyspark.sql import Window

column_prefix = "prev_Timestamp"
time_column_name = "Timestamp"
cid = 'id'


spdf= \
    spdf.withColumn(time_column_name, F.to_timestamp(F.col(time_column_name)))
w = Window.partitionBy(cid).orderBy(time_column_name)
spdf= \
    spdf\
    .withColumn(column_prefix, F.lag(spdf[time_column_name]).over(w))\
    .withColumn(column_prefix + "_DiffInSeconds", F.col(time_column_name).cast(T.LongType()) - F.col(column_prefix).cast(T.LongType()))\
    .withColumn(column_prefix + "_DiffInMinutes",F.round(F.col(column_prefix + "_DiffInSeconds")/60))\
    .withColumn(column_prefix + "_DiffInHours",F.round(F.col(column_prefix + "_DiffInSeconds")/3600))\
    .withColumn(column_prefix + "_DiffInDays",F.round(F.col(column_prefix + "_DiffInSeconds")/(24*3600)))

%%time
folder_path = '/mnt/share/data/save_folder'
spdf.write.mode('overwrite').csv(folder_path, header=True)

%%time
folder_path = '/mnt/share/data/save_folder'
spdf.coalesce(1).write.mode('overwrite').csv(folder_path, header=True)

%%time
folder_path = '/mnt/share/data/save_folder'
spdf.write.partitionBy('yyyymm').mode('overwrite').format("csv").save(folder_path, header=True)

%%time
folder_path = '/mnt/share/data/save_folder'
spdf.coalesce(1).write.partitionBy('yyyymm').mode('overwrite').format("csv").save(folder_path, header=True)

%%time
folder_path = '/mnt/share/data/save_folder'
spdf= spark.read.csv(folder_path, inferSchema=True, header=True)
# spdf.printSchema()

df= df.withColumn('label',\
    F.when(F.col('a') <= 1, 'label_1')\
     .when(F.col('a') <= 2, 'label_2')\
     .otherwise('label_3')\
)

df.groupBy('id').agg(F.sum(F.col('price')).alias('sum_of_price'))

df.groupBy('category')\
.agg({'id': 'count', 'price': 'sum'})\
.withColumnRenamed('sum(price)', 'sum_of_price')\
.withColumnRenamed('count(yid)', 'count')\
.select('category', 'sum_of_price')

spdf.select([F.count(F.when(F.isnull(F.col(c)), c)).alias(c) for c in spdf.columns]).show()

spdf = spdf.\
    .withColumn('string_column', F.lit(None).cast(T.StringType()))\
    .withColumn('int_columns', F.lit(None).cast(T.IntegerType()))\

spdf = spdf\
    .withColumn('substr', F.substring(F.col('string_column'),1,1))

spdf.write.mode("overwrite").parquet("s3a://bucket/spdf.parquet")

test = spark.read.parquet("s3a://bucket/spdf.parquet",inferSchema=True, header=True)

spdf.write.option("quoteAll", "true").mode("overwrite").csv("s3a://bucket/spdf.split.csv", header=True)

test = spark.read.csv("s3a://bucket/spdf.split.csv", header=True, inferSchema=True)

spdf.coalesce(1).write.option("quoteAll", "true").option('sep', '\t').mode("overwrite").csv("s3a://bucket/spdf.tsv", header=True)

test = spark.read.csv("s3a://bucket/spdf.tsv", sep=r'\t', header=True, inferSchema=True)

columns_to_drop=['xxx','yyy','zzz']
spdf.drop(*columns_to_drop)

from pyspark.ml.feature import Imputer

imputer = Imputer(
    inputCols=fill_by_average_columns, 
    outputCols=fill_by_average_columns
    ).setStrategy("median")
spdf_f1f2list_feature = imputer.fit(spdf_f1f2list_feature).transform(spdf_f1f2list_feature)

spdf.withColumn('total', sum(spdf[col] for col in spdf.columns)).select('total').show()

from pyspark.sql import functions as F
from pyspark.sql.window import Window as W

df = spark.createDataFrame(
  [['a', '201906'], ['a', '201907'], ['b', '201906'], ['b', '201907'], ['c', '201907']],
  ['name', 'date']
)
df.show()

result_df = (
  df
    .withColumn('max_date', F.max('date').over(W.partitionBy()))
    .filter(f.col('date') == F.col('max_date'))
    .drop('max_date')
)
result_df.show()

df = spark.createDataFrame(
    pd.DataFrame(
        data=[
            [["hello"], ["world"]],
            [["hello", "world"], ["world"]],
            [["sample", "overflow", "text"], ["sample", "text"]],
        ],
        columns=["A", "B"],
    )
)
df.show(truncate=False)

# array_exceptでAにあってBにない要素を抽出
df2 = dff.withColumn('difference', F.array_except('A', 'B'))
df2.show()
# array_intersectで共通する要素抽出
df3 = dff.withColumn('difference', F.array_intersect('A', 'B')).withColumn(
    'size_difference', F.size('difference')
)
dff3.show()

# json_columnが[{key1:value01, ket2:value02}, {key1:value11, ket2:value12},,, ]である場合
    .withColumn(
        # array of dictのstructの列を追加
        'from_json_column',
        F.from_json(
            'json_column',
            T.ArrayType(
                T.StructType(
                    [
                        T.StructField("key1", T.IntegerType()),
                        T.StructField("key2", T.StringType()),
                    ]
                )
            ),
        ),
    )