How read parquet files and only keep files that contain some columns

Viewed 311

I have a bunch of parquet files in an S3 bucket. The files contain different columns. I would like to read the files and create a datframe only with the files that contain some columns.

for example: let's say I have three columns "name", "city" and "years". Some of my files only contain, "name and "city", other contains "name","city" and "year". How can I create a dataframe by excluding the files that do not contain the column "year". I am working with spark and scala. any help is welcome.

1 Answers

How can I create a dataframe by excluding the files that do not contain the column "year".

First off I would advise restructuring bucket to separate these files based on their schema, or better yet have a process which transforms these "raw" files into a common schema that would be more easy to work with.


Working with what you have, starting with some parquet files:

val df1 = List(
 ("a", "b", "c")
).toDF("name", "city", "years")

df1.write.parquet("s3://{bucket}/test/a.parquet")

val df2 = List(
 ("aa", "bb")
).toDF("name", "city")

df2.write.parquet("s3://{bucket}/test/b.parquet")

val df3 = List(
 ("aaa", "bbb", "ccc")
).toDF("name", "city", "year")

df3.write.parquet("s3://{bucket}/test/c.parquet")

We can do:

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.{col, input_file_name}

// determine the s3 paths of all the parquet files, to be read independantly
val s3Paths = spark
  .read
  .schema(schema)
  .parquet("s3://{bucket}/test/")
  .withColumn("filename", input_file_name())
  .select("filename")
  .collect()
  .map(_(0).toString)

// individual DataFrames of each parquet file, where the `year` is not present
val dfs = s3Paths.flatMap {
  path => 
    val df = spark.read.parquet(path)
    if (!df.columns.contains("year")) {
      List(df)
    } else {
      List.empty[DataFrame]
    }
}

// take the first DataFrame, and the rest
val (firstDFs, otherDFs) = (dfs.head, dfs.tail)

// combine all of the DataFrame, unioning the rows
otherDFs.foldLeft(firstDFs) {
  case (acc, df) => acc.unionByName(df, allowMissingColumns = true)
}.show()

Note:

In the above example, when creating the test data:

df1.write.parquet("s3://{bucket}/test/a.parquet")

will have created the a file in s3, for example:

s3://{bucket}/test/a.parquet/part-0000-blah.parquet

for the purpose of creating this example I moved the parquet files up a level into the test/ path.

Related