pyspark list iterate to variable

As it pertains to:

I'm stuck on keeping the {names : new_counts}

,at the end of your for loop you may use

new_counts_list[i]=i_new

instead of

new_counts_list.append(i)

assuming that you change how new_counts_list is initialized. i.e. initialized as a dict (new_counts_list={}) instead of a list.

You also seem to be hardcoding the literal value 'i' which is a string instead of using the variable i (i.e.) without the quotes in your proposed solution. Your updated solution may look like

new_counts_list={}
for i in mylist:
    i = df.filter(df.S_ID == i).count()
    
    i_current = spark.read.format(SNOWFLAKE_SOURCE_NAME).options(**sfOptions).option("query", "select R_ID FROM mytable WHERE S_ID = '{0}'".format(i)).load()
    i_current = i_current.count()
    
    i_new = i - i_current
    i_new = str(i_new)

    new_counts_list[i]=i_new

On another note, while your approach, i.e. sequentially looping through each S_ID in my list and running the operations i.e.

1. Running the action collect to pull all the S_ID to your driver node from your initial dataframe df into a list mylist
2. Separately counting the number of occurrences of S_ID in your initial dataframe then executing another potentially expensive (IO reads/network communication/shuffles) collect()
3. Creating another dataframe with spark.read.format(SNOWFLAKE_SOURCE_NAME) that will load all records filtered by each S_ID into memory before executing a count
4. Finding the difference between the initial dataframe and the snowflake source

will work, it is expensive in IO reads and based on your cluster/setup, potentially expensive in network communication and shuffles.

You may consider using a groupby to reduce the amount of times you execute the potentially expensive collect. Furthermore, you may also join the initial dataframe to your snowflake source and let spark optimize your operations as a lazy execution plan distributed across your cluster/setup. Moreover, similar to how you are using the pushdown filter for the snowflake source, you may combine all selected S_ID in that query to allow snowflake to reduce all the desired results in one read. You would not need a loop. This could potentially look like:

Approach 1

In this approach, I will provide a pure spark solution to achieving your desired results

from pyspark.sql import functions as F
# Ask spark to select only the `S_ID` and group the data but not execute the transformation
my_exiting_counts_df = df.select('S_ID').groupBy('S_ID').count()

# Ask spark to select only the `S_ID` counts from the snowflake source
current_counts_df = (
    spark.read
         .format(SNOWFLAKE_SOURCE_NAME)
         .options(**sfOptions)
         .option("query", "select R_ID, COUNT(1) as cnt FROM mytable GROUP BY R_ID")
)

# Join both datasets which will filter to only selected `S_ID`
# and determine the differences between the existing and current counts
results_df = (
    my_exiting_counts_df.alias("existing")
                        .join(
                            current_counts_df.alias("current"),
                            F.col("S_ID")=F.col("R_ID"),
                            "inner"
                        )
                        .selectExpr(
                            "S_ID",
                            "count - cnt as actual_count"
                        )
)

# Execute the above transformations with `collect` and 
# Convert the dictionary values in the list above to your desired final dictionary
new_counts = {}
for row in results_df.collect():
    new_counts[row['S_ID']]=row['actual_count']

# your desired results are in `new_counts`

Approach 2

In this approach, I will collect the results of the group by and then use this to optimize the push down queries to the snowflake schema to return the desired results.

my_list_counts = df.select('S_ID').groupBy('S_ID').count()
selected_sids = []
case_expression = ""
for row in my_list_counts:
    selected_sids.append(row['S_ID'])
    case_expression = case_expression + " CASE WHEN R_ID='{0}' THEN {0} ".format(
        row['S_ID'],
        row['count']
    )

# The above has a table with columns `S_ID` and `count` where the 
# latter is the number of occurrences of `S_ID` in the dataset `df`

snowflake_push_down_query="""
SELECT
    R_ID AS S_ID
    ((CASE
          {0}
      END) - cnt) as actual_count
FROM (
    SELECT
        R_ID,
        COUNT(1) AS cnt
    FROM 
        mytable 
    WHERE
        R_ID IN ('{1}')
    GROUP BY 
        R_ID
) t
""".format(
   case_expression,
    "','".join(selected_sids)
)

results_df = (
    spark.read
         .format(SNOWFLAKE_SOURCE_NAME)
         .options(**sfOptions)
         .option("query", snowflake_push_down_query)
)

# Execute the above transformations with `collect` and 
# Convert the dictionary values in the list above to your desired final dictionary
new_counts = {}
for row in results_df.collect():
    new_counts[row['S_ID']]=row['actual_count']

# your desired results are in `new_counts`

Let me know if this works for you.