Repository: ray-project/ray
File: python/ray/data/preprocessors/scaler.py (Python)

Code before:
        # Read all input columns first to avoid reading modified data when
        # output_columns[i] == columns[j] for i < j
        input_columns = [table.column(input_col) for input_col in self.columns]

        for input_col, output_col, column in zip(
            self.columns, self.output_columns, input_columns
        ):
            s_mean = self.stats_[f"mean({input_col})"]
            s_std = self.stats_[f"std({input_col})"]

            if s_std is None or s_mean is None:
                # Return column filled with nulls, preserving original column type
                null_array = pa.nulls(len(column), type=column.type)
                table = BlockAccessor.for_block(table).upsert_column(
                    output_col, null_array
                )
                continue

            # Handle division by zero
            if s_std == 0:
                s_std = 1

            # Apply standard scaling: (x - mean) / std
            # Use float scalars to ensure float division
            scaled_column = pc.divide(
                pc.subtract(column, pa.scalar(float(s_mean))), pa.scalar(float(s_std))
            )

            table = BlockAccessor.for_block(table).upsert_column(
                output_col, scaled_column
            )

        return table

    @classmethod
    @DeveloperAPI
    def preferred_batch_format(cls) -> BatchFormat:
        return BatchFormat.ARROW

    def _get_serializable_fields(self) -> Dict[str, Any]:
        return {
            "columns": self.columns,
            "output_columns": self.output_columns,
            "_fitted": getattr(self, "_fitted", None),
        }

    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
        # required fields
        self.columns = fields["columns"]
        self.output_columns = fields["output_columns"]
        # optional fields

Review comment (suggestion):
```
def _scale_column(
    column: "pyarrow.Array",
    mean: Optional[float],
    std: Optional[float],
) -> "pyarrow.Array":
    """Apply standard scaling: (x - mean) / std."""
    if mean is None or std is None:
        return pa.nulls(len(column), type=column.type)
    
    # Avoid division by zero
    std = std if std != 0 else 1.0
    
    return pc.divide(
        pc.subtract(column, pa.scalar(float(mean))),
        pa.scalar(float(std)),
    )


# Read all input columns first to avoid reading modified data when
# output_columns[i] == columns[j] for i < j
input_columns = [table.column(col) for col in self.columns]

for input_col, output_col, column in zip(
    self.columns, self.output_columns, input_columns
):
    scaled = _scale_column(
        column,
        mean=self.stats_[f"mean({input_col})"],
        std=self.stats_[f"std({input_col})"],
    )
    table = BlockAccessor.for_block(table).upsert_column(output_col, scaled)

return table
```

Code after:
        if std == 0:
            std = 1

        return pc.divide(
            pc.subtract(column, pa.scalar(float(mean))), pa.scalar(float(std))
        )

    def _transform_arrow(self, table: pa.Table) -> pa.Table:
        """Transform using fast native PyArrow operations."""
        # Read all input columns first to avoid reading modified data when
        # output_columns[i] == columns[j] for i < j
        input_columns = [table.column(input_col) for input_col in self.columns]

        for input_col, output_col, column in zip(
            self.columns, self.output_columns, input_columns
        ):
            s_mean = self.stats_[f"mean({input_col})"]
            s_std = self.stats_[f"std({input_col})"]

            if s_std is None or s_mean is None:
                # Return column filled with nulls, preserving original column type
                null_array = pa.nulls(len(column), type=column.type)
                table = BlockAccessor.for_block(table).upsert_column(
                    output_col, null_array
                )
                continue

            scaled_column = self._scale_column(column, s_mean, s_std)

            table = BlockAccessor.for_block(table).upsert_column(
                output_col, scaled_column
            )

        return table

    @classmethod
    @DeveloperAPI
    def preferred_batch_format(cls) -> BatchFormat:
        return BatchFormat.ARROW

    def _get_serializable_fields(self) -> Dict[str, Any]:
        return {
            "columns": self.columns,
            "output_columns": self.output_columns,
            "_fitted": getattr(self, "_fitted", None),
        }

    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
        # required fields
        self.columns = fields["columns"]
        self.output_columns = fields["output_columns"]