[abhishekverma-ray]

Summary

Turns on deferred, batched Chain fitting by default (added behind a flag in #66091), and moves the three remaining scalers onto aggregation-based fitting so they batch inside chains too.

Before

• DataContext.enable_aggregation_based_preprocessors defaulted to False, so Chain.fit() fitted members one by one: one full dataset scan per fittable member.
• Even with the flag on, MinMaxScaler, MaxAbsScaler and RobustScaler ran their own Dataset.aggregate() inside _fit, so any chain containing them fell back to the sequential path.

After

• The flag defaults to True: Chain.fit() registers members' aggregations and computes them at the first transform() (or transform_batch() / serialization), one scan per dependency level. Opt out with RAY_DATA_ENABLE_AGGREGATION_BASED_PREPROCESSORS=0 or the DataContext setting. The flip is isolated in this PR so it's trivially revertible.
• MinMaxScaler (Min/Max), MaxAbsScaler (AbsMax) and RobustScaler register aggregators on the stat computation plan and support deferred fit.
• RobustScaler now builds one single-quantile ApproximateQuantile sketch per statistic (low_quantile, median, high_quantile) instead of one sketch queried at three points. That's a little more sketch work, in exchange for keeping every stats_ key unchanged.
• All stats_ keys and values are unchanged. KBins discretizers keep their eager fit (they post-process in _fit_execute); chains containing them, or any member without _supports_deferred_fit, still fit sequentially.

flowchart TD
    subgraph Before
        subgraph BS[Subsystem: preprocessors]
            subgraph BC["Component: Chain.fit() (flag default off)"]
                B1[MinMaxScaler._fit: own aggregate scan] --> B2[MaxAbsScaler._fit: own aggregate scan]
                B2 --> B3[RobustScaler._fit: own aggregate scan]
                B3 --> B4[StandardScaler._fit: own aggregate scan]
            end
        end
    end
    subgraph After
        subgraph AS[Subsystem: preprocessors]
            subgraph AC["Component: Chain.fit() (flag default on)"]
                A1[MinMax / MaxAbs / Robust / Standard: register aggregators on plan]
                A1 --> A2["first transform(): one Dataset.aggregate() per dependency level"]
            end
        end
    end
    classDef removed fill:#ffcdd2,stroke:#c62828,color:#000
    classDef added fill:#c8e6c9,stroke:#2e7d32,color:#000
    classDef changed fill:#fff9c4,stroke:#f9a825,color:#000
    class B1,B2,B3,B4 removed
    class A1 changed
    class A2 added

Loading

Subsystems

Preprocessors (preprocessors/scaler.py, default in context.py); computes statistics through Shuffle / aggregation (Dataset.aggregate) with no changes there.

Related issue

RFC: #66090. Final PR of the stack #66087 → #66088 → #66091 (all merged). Not a duplicate: the only open PR touching this flag is this one.

Testing

macOS / Python 3.12, env pinned to CI versions (pandas 2.3.3, pyarrow 23.0.1, cloudpickle 3.1.1, pytest 7.4.4):

• pytest python/ray/data/tests/preprocessors/test_chain.py python/ray/data/tests/preprocessors/test_scaler.py: 64 passed, including a new test that a chain of MinMax/MaxAbs/Robust/Standard scalers on disjoint columns fits in a single aggregation query with the exact expected stats_. Existing test_scaler.py assertions pass unmodified.
• pytest python/ray/data/tests/preprocessors/ --ignore=python/ray/data/tests/preprocessors/test_torch.py (torch not installed): 442 passed, 5 skipped, 0 failed (master: 441 passed).
• pytest python/ray/data/tests/test_aggregations.py: 62 passed.
• pre-commit, pyrefly (CI file set) and the API param coverage check pass.

AI assistance was used to draft this change (Claude Code); every line has been reviewed by the submitter.

🤖 Generated with Claude Code

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[gemini-code-assist]

Code Review

This pull request enables aggregation-based preprocessors by default and implements deferred fitting support for MinMaxScaler, MaxAbsScaler, and RobustScaler by refactoring their _fit methods to register aggregators with the stat computation plan. Additionally, the test suite is updated to verify that multiple scalers in a chain can defer fitting and execute in a single aggregation query. The review feedback correctly identifies a potential late-binding issue in RobustScaler._fit, where self._quantile_precision is captured from the outer scope of a deferred lambda, and suggests binding it as a default argument to ensure safety during serialization or mutation.

[gemini-code-assist]

In RobustScaler._fit, the lambda function for aggregator_fn captures self._quantile_precision from the outer scope. Since this lambda is executed deferredly during statistics materialization (which happens later during transform), any potential mutation or serialization/deserialization of the preprocessor instance could lead to unexpected behavior or lookup issues. Capturing self._quantile_precision as a default argument (e.g., prec=self._quantile_precision) at definition time ensures that the value is safely bound when _fit is called, matching the behavior of q=quantile and name=stat_name.

Suggested change

	self._stat_computation_plan.add_aggregator(
	aggregator_fn=lambda col, q=quantile, name=stat_name: (
	ApproximateQuantile(
	on=col,
	quantiles=[q],
	quantile_precision=self._quantile_precision,
	alias_name=f"{name}({col})",
	)
	),
	post_process_fn=unpack_quantile,
	columns=self._columns,
	)
	self._stat_computation_plan.add_aggregator(
	aggregator_fn=lambda col, q=quantile, name=stat_name, prec=self._quantile_precision: (
	ApproximateQuantile(
	on=col,
	quantiles=[q],
	quantile_precision=prec,
	alias_name=f"{name}({col})",
	)
	),
	post_process_fn=unpack_quantile,
	columns=self._columns,
	)