[Data] Add serialization framework for preprocessors

[cem-anyscale]

Description

This commit introduces a new serialization system for Ray Data preprocessors that improves maintainability, extensibility, and backward compatibility.

Key changes:

1. New serialization infrastructure:

   • Add serialization_handlers.py with factory pattern for format handling
   • Implement CloudPickleSerializationHandler (primary format)
   • Support legacy PickleSerializationHandler for backward compatibility
   • Add format auto-detection via magic bytes (CPKL:)

2. New preprocessor base class:

   • Add SerializablePreprocessorBase abstract class
   • Define serialization interface via abstract methods:
     • _get_serializable_fields() / _set_serializable_fields()
     • _get_stats() / _set_stats()
   • Mark serialize() and deserialize() as @Final to prevent overrides

3. Preprocessor registration system:

   • Add version_support.py with @SerializablePreprocessor decorator
   • Enable versioned serialization with stable identifiers
   • Support class registration and lookup
   • Add UnknownPreprocessorError for missing types

4. Migrate preprocessors to new framework:

   • SimpleImputer
   • OrdinalEncoder
   • OneHotEncoder
   • MultiHotEncoder
   • LabelEncoder
   • Categorizer
   • StandardScaler
   • MinMaxScaler
   • MaxAbsScaler
   • RobustScaler

5. Enhanced Preprocessor base class:

   • Add get_input_columns() and get_output_columns() methods (for future use)
   • Add has_stats() (for future use)
   • Add type hints to getstate() and setstate()

6. Backward compatibility improvements to Concatenator for existing functionality:

   • Add setstate override in Concatenator for flatten field
   • Handle missing fields gracefully during deserialization

The new architecture makes it easier to:

• Add new serialization formats without modifying core logic
• Maintain backward compatibility with existing serialized data
• Handle version migrations for preprocessor schemas
• Register new preprocessors with stable identifiers

[cursor]

Bug: Test Fails Due to Hardcoded Serialization

The test_encoder_serialization_formats() method incorrectly claims to test different serialization formats. The SerializablePreprocessorBase.serialize() method is hardcoded to CloudPickle, causing the test to compare two identical CloudPickle outputs. Additionally, the serialize() docstring incorrectly lists an output_format parameter.

Additional Locations (1)

• python/ray/data/preprocessor.py#L618-L620

 

[gemini-code-assist]

Code Review

This pull request introduces a robust new serialization framework for Ray Data preprocessors, which is a significant improvement for maintainability, extensibility, and backward compatibility. The changes are well-structured, introducing a new SerializablePreprocessorBase class, a factory pattern for handling different serialization formats, and a versioned registration system for preprocessors. The migration of existing preprocessors to this new framework is well-executed. The code quality is high, and the new functionality is accompanied by a comprehensive suite of tests. My review focuses on a few minor documentation inconsistencies and a potentially misleading test case.

[gemini-code-assist]

The docstring for serialize mentions an output_format argument, but this argument is not present in the method signature. This can be confusing for developers using this API. Please remove the output_format from the Args section of the docstring.

[gemini-code-assist]

The docstring for get_handler is incomplete. The Args section is empty. Please document the parameters format_identifier, data, and **kwargs to improve clarity for developers who might use or extend this factory.

[gemini-code-assist]

This test test_encoder_serialization_formats seems to have a misleading comment and implementation. It claims to test the legacy Pickle format, but it calls fitted_encoder.serialize(), which for a SerializablePreprocessorBase subclass will always use the new CloudPickle-based serialization. The old pickle-based serialization produced a str, while the new one produces bytes.

The test is effectively re-running the CloudPickle serialization test. To properly test backward compatibility for deserialization of the legacy pickle format, you would need to either manually create a legacy-formatted serialized string or have a pre-serialized object from an older version.

Given that this test doesn't check what it claims to, it would be best to either fix it to correctly test legacy format deserialization or remove it to avoid confusion.

[dstrodtman]

stamp

[cursor]

Bug: Serialization Skips Crucial Imputer Validation

The _set_serializable_fields method doesn't validate that fill_value is not None when strategy is "constant", unlike the constructor which raises a ValueError for this invalid combination. Deserializing a SimpleImputer with strategy="constant" and fill_value=None creates an invalid object that will fail during transformation with a confusing error message.

 

[cursor]

Bug: Serialize: Unexpected Return Type

The Preprocessor.serialize() method's docstring and return type annotation claim it returns str, but SerializablePreprocessorBase.serialize() (which overrides this method for new preprocessors) returns Union[str, bytes] and actually returns bytes for CloudPickle format. This creates a breaking change where code expecting a string from serialize() will receive bytes instead, potentially causing type errors or incorrect behavior in downstream code that assumes string output.