Create N tables from a base table with table names derived from column values without a for loop

siva 0
Viewed 26

I have one big origin table which I want to split in several smaller tables based on the category.

id Category color
1 car yellow
2 plane blue
3 plane green
4 car blue
5 bus blue
5 plane yellow 
Table name: car
| id | Category | color |
| ---| -------- | ----- |
| 1  | car      | yellow|
| 4  | car      | blue  |

Table name: plane

id Category color
2 plane blue
3 plane green
5 plane yellow

Table name: bus

id Category color
5 bus blue

This is one potential solution:

df_tmp = df_transportation.groupBy("category").count()
categories = list(df_tmp["category"])
for category in categories:
   df_cat = df_transportation.filter(col("category") == "category))
   df_cat.write.mode("overwrite").format("delta").saveAsTable(f'{category}'

Is there a way to get to the same solution without a for loop?

pyspark apache-spark
1 Answers

Yes it can be done with a one liner.

# create df
df = spark.createDataFrame(sc.parallelize([
    [1, 'A', 20220722, 1],
    [1, 'A', 20220723, 1],
    [1, 'B', 1, 0],
    [2, 'B', 20220722, 1],
    [2, 'C', 0, 0],
    [2, 'B', 20220724, 3],
]),
                           ['ID', 'State', 'Time', 'Expected'])

+---+-----+--------+--------+
| ID|State|    Time|Expected|
+---+-----+--------+--------+
|  1|    A|20220722|       1|
|  1|    A|20220723|       1|
|  1|    B|       1|       0|
|  2|    B|20220722|       1|
|  2|    C|       0|       0|
|  2|    B|20220724|       3|
+---+-----+--------+--------+

Solution

l  = dict(zip(v:=sorted(set(df.select('State').rdd.flatMap(lambda x: x).collect())),[df.where(col('State')==x) for x in v] ))

l.get('B').show()

+---+-----+--------+--------+
| ID|State|    Time|Expected|
+---+-----+--------+--------+
|  1|    B|       1|       0|
|  2|    B|20220722|       1|
|  2|    B|20220724|       3|
+---+-----+--------+--------+

How it works

I create a dictionary of element and its filter in the dataframe.

Using; v:=sorted(set(df.select('State').rdd.flatMap(lambda x: x).collect())) Using the walrus operator I store the unique elements in the column State in a variable v

using; [df.where(col('State')==x) for x in v] I store filtered dataframe of each element in the main df

Using dict(zip()) I store the two outcomes above in a dictionary.

using l.get('B').show() I call a dataframe in the dict whose key is B

wwnde0
Related