Skip to main content

Spark DataFrame Serialization Internals

In Flyte, passing Spark DataFrames between tasks requires a serialization mechanism that can handle distributed data efficiently. The vik-advani-flyte-sdk-9b3ce04 codebase implements this using Parquet as the intermediate wire format, managed by the SparkToParquetEncoder and ParquetToSparkDecoder classes in flyteplugins.spark.df_transformer.

This implementation bridges the gap between Spark's distributed memory model and Flyte's StructuredDataset type system, ensuring that large datasets can be offloaded to remote storage (like S3 or GCS) and re-loaded by downstream tasks.

The Role of Parquet as a Wire Format​

The choice of Parquet for Spark serialization is driven by its columnar nature and native integration with the Spark ecosystem. In flyteplugins.spark.df_transformer, both the encoder and decoder explicitly declare PARQUET as their supported format:

class SparkToParquetEncoder(DataFrameEncoder):
    def __init__(self):
        super().__init__(python_type=pyspark.sql.DataFrame, supported_format=PARQUET)

class ParquetToSparkDecoder(DataFrameDecoder):
    def __init__(self):
        super().__init__(pyspark.sql.DataFrame, None, PARQUET)

By using Parquet, Flyte leverages Spark's ability to perform high-performance parallel I/O, which is essential when tasks are running on different physical nodes in a Kubernetes cluster.

Serialization Logic (Encoding)​

The SparkToParquetEncoder.encode method handles the transition from a live Spark DataFrame to a persistent Parquet file. This process involves several critical design decisions regarding path management and Spark configuration.

Path Resolution​

If a URI is not already associated with the DataFrame, the encoder dynamically resolves a storage path using Flyte's internal context. It prioritizes the raw_data_path from the task context, ensuring that the serialized data is stored in the correct remote location configured for the Flyte project.

if not path:
    tctx = flyte.ctx()
    if tctx is not None:
        path = tctx.raw_data_path.get_random_remote_path()

Spark Session and Hygiene​

The encoder retrieves the current SparkSession and applies a specific configuration to prevent the generation of _SUCCESS files. In a distributed workflow system like Flyte, these marker files are often redundant or can interfere with directory-based data reading in certain object stores.

ss = pyspark.sql.SparkSession.builder.getOrCreate()

# Avoid generating SUCCESS files
ss.conf.set("mapreduce.fileoutputcommitter.marksuccessfuljobs", "false")
cast(pyspark.sql.DataFrame, dataframe._raw_df).write.mode("overwrite").parquet(path=path)

The use of mode("overwrite") is a significant design choice, ensuring that retried tasks or overlapping executions do not fail due to existing data at the target URI, though it implies that the URI should be unique per execution attempt.

Deserialization Logic (Decoding)​

The ParquetToSparkDecoder.decode method is responsible for reading the Parquet data back into a pyspark.sql.DataFrame.

Column Projection Optimization​

A key feature of this implementation is the support for column projection. If the Flyte type signature for the input specifies a subset of columns, the decoder uses Spark's .select() method to read only those columns. This reduces memory usage and I/O overhead by leveraging Parquet's columnar storage.

if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
    columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
    return spark.read.parquet(path).select(*columns)
return spark.read.parquet(path)

Registration and Integration​

The Spark transformers are integrated into Flyte's TypeEngine via the DataFrameTransformerEngine. This registration is handled by register_spark_df_transformers(), which is decorated with @functools.lru_cache to ensure registration happens exactly once per process.

@functools.lru_cache(maxsize=None)
def register_spark_df_transformers():
    DataFrameTransformerEngine.register(SparkToParquetEncoder(), default_format_for_type=True)
    DataFrameTransformerEngine.register(ParquetToSparkDecoder(), default_format_for_type=True)

This function is called at the module level in flyteplugins/spark/df_transformer.py, ensuring that as soon as the Spark plugin is imported, Flyte knows how to handle pyspark.sql.DataFrame types.

Tradeoffs and Constraints​

  1. Spark Session Dependency: Both the encoder and decoder rely on pyspark.sql.SparkSession.builder.getOrCreate(). This assumes that the environment where the conversion happens (the Flyte task pod) is correctly configured with a Spark session, which is typically handled by the Flyte Spark plugin's task executor.
  2. Overwrite Semantics: The hardcoded overwrite mode in the encoder simplifies task retries but requires that the raw_data_path generates unique URIs to avoid data corruption across different task instances.
  3. Format Rigidity: While StructuredDataset supports multiple formats, this specific implementation is tightly coupled to PARQUET. Users requiring other formats (like Avro or ORC) would need to implement separate encoders/decoders or extend these classes.