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

Code before:
        self.stat_computation_plan.add_aggregator(
            aggregator_fn=lambda col: Std(col, ddof=0),
            columns=self.columns,
        )
        return self

    def _transform_pandas(self, df: pd.DataFrame):
        def column_standard_scaler(s: pd.Series):
            s_mean = self.stats_[f"mean({s.name})"]
            s_std = self.stats_[f"std({s.name})"]

            if s_std is None or s_mean is None:
                s[:] = np.nan
                return s

            # Handle division by zero.
            # TODO: extend this to handle near-zero values.
            if s_std == 0:
                s_std = 1

            return (s - s_mean) / s_std

        df[self.output_columns] = df[self.columns].transform(column_standard_scaler)
        return df

    def _set_or_append_column(
        self, table: pa.Table, column: pa.ChunkedArray, output_col: str
    ) -> pa.Table:
        column_idx = table.schema.get_field_index(output_col)
        if column_idx == -1:
            return table.append_column(output_col, column)
        else:
            return table.set_column(column_idx, output_col, column)

    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 = self._set_or_append_column(table, null_array, output_col)
                continue

Review comment (suggestion):
Since ur using this function multiple times, I suggest doing something like from this comment: https://github.com/ray-project/ray/pull/59890#discussion_r2673241444. Try and see if u can do most operations with `ArrowBlockAccessor`

Code after:
        )
        self.stat_computation_plan.add_aggregator(
            aggregator_fn=lambda col: Std(col, ddof=0),
            columns=self.columns,
        )
        return self

    def _transform_pandas(self, df: pd.DataFrame):
        def column_standard_scaler(s: pd.Series):
            s_mean = self.stats_[f"mean({s.name})"]
            s_std = self.stats_[f"std({s.name})"]

            if s_std is None or s_mean is None:
                s[:] = np.nan
                return s

            # Handle division by zero.
            # TODO: extend this to handle near-zero values.
            if s_std == 0:
                s_std = 1

            return (s - s_mean) / s_std

        df[self.output_columns] = df[self.columns].transform(column_standard_scaler)
        return df

    @staticmethod
    def _scale_column(column: pa.Array, mean: float, std: float) -> pa.Array:
        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(