Is there a way to add a column of type dictionary to a spark dataframe in pyspark?

user_5295157 0
Viewed 4689

This is how I create a dataframe with primitive data types in pyspark:

from pyspark.sql.types import StructType, StructField, DoubleType, StringType, IntegerType
fields = [StructField('column1', IntegerType(), True), StructField('column2', IntegerType(), True)]
schema = StructType(fields)
df = spark.createDataFrame([], schema)
values = [tuple([i]) +
          tuple([i])
          for i in range(3)]
df = spark.createDataFrame(values, schema)

Now, if I want to have a third column with dictionary data, eg: {"1": 1.0, "2": 2.0, "3": 3.0}, what should I do? I want to create this data frame:

+--------------------+-----------------+------------------------------+
|column1             |column2          |column3                       |
+--------------------+-----------------+------------------------------+
|1                   |1                |{"1": 1.0, "2": 1.0, "3": 1.0}|
+--------------------+-----------------+------------------------------+
|2                   |2                |{"1": 2.0, "2": 2.0, "3": 2.0}|
+--------------------+-----------------+------------------------------+
|3                   |3                |{"1": 3.0, "2": 3.0, "3": 3.0}|
+--------------------+-----------------+------------------------------+

There is a MapType that seems to be helpful, but I can't figure out how to use it?

And assuming the data frame is created, how to filter it based on the third column, given a dict to select the rows of the data frame that have that dict value?

pyspark apache-spark-sql python
2 Answers

Example how to create:

from pyspark.sql.types import MapType, IntegerType, DoubleType, StringType, StructType, StructField
import pyspark.sql.functions as f

schema = StructType([
            StructField('column1', IntegerType()),
            StructField('column2', IntegerType()),
            StructField('column3', MapType(StringType(), DoubleType()))])

data = [(1, 2, {'a':3.5, 'b':4.2}), (4, 8, {'b':3.7, 'e':4.9})]
df = spark.createDataFrame(data, schema=schema)
df.show()

Output:

+-------+-------+--------------------+
|column1|column2|             column3|
+-------+-------+--------------------+
|      1|      2|[a -> 3.5, b -> 4.2]|
|      4|      8|[e -> 4.9, b -> 3.7]|
+-------+-------+--------------------+

Example on how to filter DataFrame only leaving elements which have a certain key (assuming you don't have null values in the map and your Spark version is 2.4+ cause early versions don't have element_at):

filtered_df = df.where(f.element_at(df.column3, 'a').isNotNull())

Output:

+-------+-------+--------------------+
|column1|column2|             column3|
+-------+-------+--------------------+
|      1|      2|[a -> 3.5, b -> 4.2]|
+-------+-------+--------------------+

I might have misunderstood your question - if your intention is to only leave rows where map column equal to a specific dictionary you have it is a little bit more tricky. As far as I know Spark doesn't have comparison operation on dictionary types (it is somewhat unusual operation). There is a way to implement it using udf, which will be not very efficient. The code for that might look like this:

from pyspark.sql.types import MapType, IntegerType, DoubleType, StringType, StructType, StructField, BooleanType
my_dict = {'b':2.7, 'e':4.9}

from pyspark.sql.functions import udf
def map_equality_comparer(my_dict):
    @udf(BooleanType())
    def comparer(m):
        if len(m) != len(my_dict): return False
        for k, v in m.items():
            if my_dict.get(k) != v: return False
        return True
    return comparer

filtered_df = df.where(map_equality_comparer(my_dict)(df.column3))
filtered_df.show()

If this is too slow for you you might consider creating a canonical representation of your Dictionaries and comparing those (e.g. converting dictionaries to sorted arrays of key value pairs and filtering based on equality of these arrays).

user_118396860

You can do this without UDF using, create_map(spark2.0+), you can create a new column of MapType, using ur dictionary to be used, and then use filter as shown below to get all rows where my_dict is equal to column3.

from pyspark.sql import functions as F


df.show() #sample dataframe
my_dict = {'b':3.7, 'e':4.9} #dictionary to filter with

#+-------+-------+--------------------+
#|column1|column2|             column3|
#+-------+-------+--------------------+
#|      1|      2|[a -> 3.5, b -> 4.2]|
#|      4|      8|[e -> 4.9, b -> 3.7]|
#+-------+-------+--------------------+

from pyspark.sql import functions as F

df.withColumn("map", F.create_map(*[item for sublist in [[F.lit(x),F.lit(y)]\
                                for x,y in my_dict.items()] for item in sublist]))\
  .filter(' and '.join(["column3.{0}=map.{0}".format(x) for x in my_dict.keys()])+\
          ' and size(column3)=size(map)').drop("map").show()

#+-------+-------+--------------------+
#|column1|column2|             column3|
#+-------+-------+--------------------+
#|      4|      8|[e -> 4.9, b -> 3.7]|
#+-------+-------+--------------------+
user_101073890
Related