How can I drop columns that are not part of a schema safely?

user_1185587 0
Viewed 267

I have a dataframe that has columns: [A, B, ... M] and a schema that only requires few of the columns from the dataframe:

StructType([StructField(A, StringType(), False),
            StructField(C, StringType(), True),
            StructField(K, StringType(), True)])

Since I do not own the schema, I do not have a list of all columns already available. Is there a way to safely and efficiently select columns as per the schema?

pyspark apache-spark-sql apache-spark
2 Answers

A "safer" way is to check whether each column in the provided schema is present in the dataframe, and select all those that are present using a list comprehension:

schema = StructType([StructField("A", StringType(), False),
                     StructField("C", StringType(), True),
                     StructField("K", StringType(), True)])

df2 = df.select(*[c for c in schema.names if c in df.columns])
user_141657300

lets say schema you have is

schema = StructType([StructField(A, StringType(), False),
            StructField(C, StringType(), True),
            StructField(K, StringType(), True)])

to select only these columns from dataframe you can do something like this

df.select([col for col in schema])
user_70762750
Related