chore: import upstream snapshot with attribution

python/ray/data/_internal/random_config.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+from ray.data.context import DataContext
+from ray.util.annotations import DeveloperAPI
+
+# NumPy's RandomState/seed range is [0, 2**32 - 1].
+NUMPY_RNG_SEED_MAX = 2**32
+
+
+@dataclass(frozen=True)
+class SeedTuple:
+    """A seed for random number generation, optionally including execution index for reseeding.
+
+    Args:
+        seed: The base seed.
+        execution_idx: The execution index. If None, the seed is not reseeded after execution.
+
+    """
+
+    seed: int
+    execution_idx: Optional[int] = None
+
+    def to_rng_args(self, task_idx: int) -> tuple[int, ...]:
+        """Return seed parts for np.random.default_rng((task_idx, ...)).
+
+        Args:
+            task_idx: The task index.
+
+        Returns:
+            A tuple of seed parts.
+        """
+        if self.execution_idx is None:
+            return (task_idx, self.seed)
+        return (task_idx, self.execution_idx, self.seed)
+
+
+@DeveloperAPI
+@dataclass
+class RandomSeedConfig:
+    """This configuration object controls the random seed behavior for operations
+    such as :meth:`~Dataset.random_shuffle`, :meth:`~Dataset.randomize_block_order`,
+    and :meth:`~Dataset.random_sample`. The random seed behavior is determined by
+    the combination of the base seed ``seed`` and the ``reseed_after_execution`` parameter:
+
+    - If ``seed`` is None, the random seed is always None (non-deterministic shuffling).
+    - If ``seed`` is not None and ``reseed_after_execution`` is False, the base seed is
+      used as the random seed for each execution.
+    - If ``seed`` is not None and ``reseed_after_execution`` is True, the base seed is
+      combined with the (incremental) execution index ``execution_idx`` to produce a
+      different random seed tuple for each execution.
+
+    .. note::
+        Even if you provided a seed, you might still observe a non-deterministic row
+        order. This is because tasks are executed in parallel and their completion
+        order might vary. If you need to preserve the order of rows, set
+        ``DataContext.get_current().execution_options.preserve_order``.
+
+    Args:
+        seed: An optional integer base seed. If None, the operation is
+            non-deterministic. If provided, the operation is deterministic based on the base
+            seed and the ``reseed_after_execution`` parameter.
+        reseed_after_execution: If True, the random seed considers both ``seed`` and
+            ``execution_idx``, resulting in different shuffling orders across executions.
+            If False, the base seed is used as the random seed for each execution, resulting in the same
+            shuffling order across executions. Only takes effect when a base seed is provided.
+            Defaults to True.
+        use_timestamp_as_default: When enabled, it supports a legacy behavior that relies on
+            the timestamp as the default seed. This parameter is only used when the base seed
+            is None. Defaults to False. See ``get_single_integer_random_seed`` for more details.
+
+    """  # noqa: E501
+
+    seed: Optional[int] = None
+    reseed_after_execution: bool = True
+    use_timestamp_as_default: bool = False
+
+    def __post_init__(self):
+        """Ensure that the seed is either None or an integer."""
+        if self.seed is not None and not isinstance(self.seed, int):
+            raise ValueError("Seed must be an integer or None.")
+
+    def get_seed_tuple(
+        self,
+        *,
+        data_context: DataContext,
+    ) -> SeedTuple | None:
+        """Return a seed for random number generation.
+
+        Args:
+            data_context: A DataContext object for extracting the
+                execution index.
+
+        Returns:
+            A SeedTuple, or None for non-deterministic behavior.
+        """
+        if self.seed is None:
+            return None
+        elif self.reseed_after_execution:
+            return SeedTuple(
+                seed=self.seed,
+                execution_idx=data_context._execution_idx,
+            )
+        else:
+            return SeedTuple(seed=self.seed, execution_idx=None)
+
+    @classmethod
+    def create_seed_config(
+        cls,
+        seed: int | RandomSeedConfig | None,
+        *,
+        use_timestamp_as_default: bool = False,
+    ) -> RandomSeedConfig:
+        """Create a ``RandomSeedConfig`` object from the ``seed`` argument in Ray Data public random APIs.
+
+        This is a helper function that converts an integer seed into a ``RandomSeedConfig`` object. In this case,
+        ``reseed_after_execution`` is set to False, which matches the standard behavior for a data pipeline
+        (i.e., same seed across executions). To override this behavior, use a ``RandomSeedConfig`` object
+        directly.
+
+        Args:
+            seed: This optional argument can be an integer or an existing ``RandomSeedConfig`` object.
+            use_timestamp_as_default: If True, a timestamp-based seed is used when
+                ``seed`` is None. This pins the seed at plan time so that task retries
+                produce identical output. When ``seed`` is an existing
+                ``RandomSeedConfig``, this value overrides its
+                ``use_timestamp_as_default`` field.
+
+        Returns:
+            A ``RandomSeedConfig`` object.
+        """
+        if seed is None or isinstance(seed, int):
+            return RandomSeedConfig(
+                seed=seed,
+                reseed_after_execution=False,
+                use_timestamp_as_default=use_timestamp_as_default,
+            )
+        elif isinstance(seed, RandomSeedConfig):
+            return cls(
+                seed=seed.seed,
+                reseed_after_execution=seed.reseed_after_execution,
+                use_timestamp_as_default=use_timestamp_as_default,
+            )
+
+        raise ValueError(f"Invalid seed type: {type(seed)}")
+
+
+def get_timestamp_seed() -> int:
+    """Returns a timestamp-based seed."""
+    return time.time_ns() % NUMPY_RNG_SEED_MAX
+
+
+def get_single_integer_random_seed(
+    seed_config: RandomSeedConfig,
+    data_context: DataContext,
+) -> Optional[int]:
+    """Returns a single integer seed based on the ``RandomSeedConfig`` object. This is useful
+    if the caller expects a single integer to seed the RNG. When the base seed is None,
+    and ``seed_config.use_timestamp_as_default`` is True, a timestamp is used as the seed.
+    Otherwise, the seed tuple is hashed to produce a single integer seed.
+
+    Args:
+        seed_config: The ``RandomSeedConfig`` object.
+        data_context: The ``DataContext`` object used to generate the seed tuple.
+
+    Returns:
+        A single integer random seed, or None for non-deterministic behavior.
+    """
+    seed_result = seed_config.get_seed_tuple(data_context=data_context)
+
+    if seed_result is None:
+        # This is a legacy behavior for some random operations.
+        return get_timestamp_seed() if seed_config.use_timestamp_as_default else None
+    elif seed_result.execution_idx is None:
+        return seed_result.seed
+
+    # The modulo is only needed because some random implementations are using the
+    # older type RandomState or np.random.seed(). Otherwise, the seed can be
+    # as large as 128-bit integer. See
+    # https://blog.scientific-python.org/numpy/numpy-rng/
+    return hash(seed_result) % NUMPY_RNG_SEED_MAX