Binary check code in pyspark I am trying to find out if a column is binary or not. If a column is only having 1 or 0 then I am flagging it as binary, else non binary. Since I have a database background, I tried to achieve it through a SQL like statement. But when the file is big, this code is not giving me good performance.

Could you please suggest how I can improve this code:

input_data=spark.read.csv("/tmp/sample.csv", inferSchema=True,header=True)
input_data.createOrReplaceTempView("input_data")
totcount=input_data.count()

from pyspark.sql.types import StructType,StructField,StringType
profSchema = StructType([ StructField("column_name", IntegerType(), True)\
                   ,StructField("binary_type", StringType(), True)])
fin_df=spark.createDataFrame([],schema=profSchema) ##create null df

for colname in input_data.columns:
    query="select {d_name} as column_name,case when sum(case when {f_name} in ( 1,0) then 1 else 0 end) == {tot_cnt} then 'binary' else 'nonbinary'\
    end as binary_stat from input_data".format(f_name=colname, d_name="'"+str(colname)+"'",tot_cnt=totcount)
    join_df=spark.sql(query)
    fin_df=fin_df.union(join_df)
Spark dataframes (and columns) have a distinct method, which you can use to get all values in that column. Getting the actual values out is a bit more complicated and taken from this answer to a similar question on StackOverflow:

from pyspark.sql import SparkSession

def get_distinct_values(data, column):
    return {x[column] for x in data.select(column).distinct().collect()}

spark = SparkSession \
    .builder \
    .appName("Python Spark SQL basic example") \
    .config("spark.some.config.option", "some-value") \
    .getOrCreate()

input_data = spark.read.csv("/tmp/sample.csv", inferSchema=True, header=True)
input_data.createOrReplaceTempView("input_data")

print({c: get_distinct_values(input_data, c) == {True, False}
       for c in input_data.columns})
# {'category': False, 'logged_in': True, 'gid': False, 'pol_id': False, 'subcategory': False}


I don't know enough about spark to know how you would cast this back into a spark dataframe, but this should get you at least halfway there and be a bit faster, since it can do the fastest implementation it can to reduce the values to sets.



# /tmp/sample.csv
pol_id,gid,category,subcategory,logged_in
1,1,A,a,1
2,1,A,b,0
1,2,B,b,1
2,2,B,a,0