Error when accessing rdd of DataFrame

pablobd
Contributor II

‎01-24-2024 02:18 PM

I need to run this kind of code:

 

from pyspark.sql import SparkSession
import pandas as pd

# Create a Spark session
spark = SparkSession.builder.appName("example").getOrCreate()

# Sample data
data = [("Alice", 1), ("Bob", 2), ("Charlie", 3), ("David", 4), ("Eva", 5), ("Frank", 6)]
columns = ["Name", "Value"]
df = spark.createDataFrame(data, columns)

# Function to be applied on each partition
def decrypt_on_partition(iterator):
    # Convert the iterator to a Pandas DataFrame
    partition_df = pd.DataFrame(list(iterator), columns=columns)

    # decryption client and cacher
    KEY_ARN = _get_decrypt_key()
    key_provider = aws_encryption_sdk.StrictAwsKmsMasterKeyProvider(key_ids=[KEY_ARN])
    cache = aws_encryption_sdk.LocalCryptoMaterialsCache(LOCAL_CACHE)
    caching_cmm = aws_encryption_sdk.CachingCryptoMaterialsManager(
        master_key_provider=key_provider,
        cache=cache,
        max_age=MAX_AGE_IN_SECONDS,
        max_messages_encrypted=MAX_ENTRY_MESSAGES,
    )
    client = aws_encryption_sdk.EncryptionSDKClient(
        commitment_policy=CommitmentPolicy.REQUIRE_ENCRYPT_REQUIRE_DECRYPT
    )
    
    # Add an additional column to the Pandas DataFrame
    partition_df = partition_df.assign(
        bookingDict=partition_df["data"].apply(
            lambda x: _decrypt_string(x, client, caching_cmm)
        )
    
    # Convert the Pandas DataFrame back to an iterator
    for row in partition_df.itertuples(index=False):
        yield tuple(row)
        
# Apply the function to each partition
result_rdd = df.rdd.mapPartitions(decrypt_on_partition)

# Collect the results back to the driver node (for demonstration purposes)
result_list = result_rdd.collect()

# Print the result
print(result_list)

# Stop the Spark session
spark.stop()

 

Why?

Because with this code I send the caching service and encryption client to each node and run it on each partition of the DataFrame. Also, the decryption service can only be accessed from a Spark cluster with the  appropriate AWS IAM role associated. And the cacher used by the decryption service (to speed up decryption and save on costs) can't be run in Spark (that's why I convert the stream of data to a pandas.DataFrame).
 
Benefits:
  1. It runs 4 times faster (4 nodes in this cluster). All integ data decrypted in 30 seconds.
  2. I should be able to decrypt the data from my local computer if connected to the appropriate cluster and develop the ML Model locally.

When running it in a Databricks notebook it works, however, locally I get this error:

[NOT_IMPLEMENTED] rdd() is not implemented.

I also tried to use the spark.DataFrame api without success, doing this,

 

df = df.foreachPartition(decrypt_on_partition)

 

Instead of,

 

result_rdd = df.rdd.mapPartitions(decrypt_on_partition)
result_list = result_rdd.collect()

 

Thanks a lot

0 Kudos