How to include partitioned column in pyspark dataframe read method

0
votes

I am writing Avro file-based from a parquet file. I have read the file as below:

Reading data

dfParquet = spark.read.format("parquet").option("mode", "FAILFAST")
    .load("/Users/rashmik/flight-time.parquet")

Writing data

I have written the file in Avro format as below:

dfParquetRePartitioned.write \
    .format("avro") \
    .mode("overwrite") \
    .option("path", "datasink/avro") \
    .partitionBy("OP_CARRIER") \
    .option("maxRecordsPerFile", 100000) \
    .save()

As expected, I got data partitioned by OP_CARRIER.

Reading Avro partitioned data from a specific partition

In another job, I need to read data from the output of the above job, i.e. from datasink/avro directory. I am using the below code to read from datasink/avro

dfAvro = spark.read.format("avro") \
    .option("mode","FAILFAST") \
    .load("datasink/avro/OP_CARRIER=AA")

It reads data successfully, but as expected OP_CARRIER column is not available in dfAvro dataframe as it is a partition column of the first job. Now my requirement is to include OP_CARRIER field also in 2nd dataframe i.e. in dfAvro. Could somebody help me with this?

I am referring documentation from the spark document, but I am not able to locate the relevant information. Any pointer will be very helpful.

1
apache-sparkpysparkapache-spark-sqlpyspark-dataframes
.load("datasink/avro") - Lamanus

1 Answers

0
votes

You replicate the same column value with a different alias.

dfParquetRePartitioned.withColumn("OP_CARRIER_1", lit(df.OP_CARRIER)) \
.write \
.format("avro") \
.mode("overwrite") \
.option("path", "datasink/avro") \
.partitionBy("OP_CARRIER") \
.option("maxRecordsPerFile", 100000) \
.save()

This would give you what you wanted. But with a different alias. Or you can also do it during reading. If location is dynamic then you can easily append the column.

path = "datasink/avro/OP_CARRIER=AA"
newcol = path.split("/")[-1].split("=") 
dfAvro = spark.read.format("avro") \
.option("mode","FAILFAST") \
.load(path).withColumn(newcol[0], lit(newcol[1]))

If the value is static its way more easy to add it during the data read.