datachain 0.37.7__tar.gz → 0.37.8__tar.gz

{datachain-0.37.7 → datachain-0.37.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.37.7
+Version: 0.37.8
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.37.7 → datachain-0.37.8}/src/datachain/lib/dc/datachain.py

@@ -52,7 +52,11 @@ from datachain.lib.udf_signature import UdfSignature
 from datachain.lib.utils import DataChainColumnError, DataChainParamsError
 from datachain.project import Project
 from datachain.query import Session
-from datachain.query.dataset import DatasetQuery, PartitionByType
+from datachain.query.dataset import (
+    DatasetQuery,
+    PartitionByType,
+    RegenerateSystemColumns,
+)
 from datachain.query.schema import DEFAULT_DELIMITER, Column
 from datachain.sql.functions import path as pathfunc
 from datachain.utils import batched_it, env2bool, inside_notebook, row_to_nested_dict
@@ -2740,8 +2744,20 @@ class DataChain:
         )
 
     def shuffle(self) -> "Self":
-        """Shuffle the rows of the chain deterministically."""
-        return self.order_by("sys.rand")
+        """Shuffle rows with a best-effort deterministic ordering.
+
+        This produces repeatable shuffles. Merge and union operations can
+        lead to non-deterministic results. Use order by or save a dataset
+        afterward to guarantee the same result.
+        """
+        query = self._query.clone(new_table=False)
+        query.steps.append(RegenerateSystemColumns(self._query.catalog))
+
+        chain = self._evolve(
+            query=query,
+            signal_schema=SignalSchema({"sys": Sys}) | self.signals_schema,
+        )
+        return chain.order_by("sys.rand")
 
     def sample(self, n: int) -> "Self":
         """Return a random sample from the chain.

{datachain-0.37.7 → datachain-0.37.8}/src/datachain/query/dataset.py

@@ -786,10 +786,31 @@ class SQLClause(Step, ABC):
         return tuple(c.get_column() if isinstance(c, Function) else c for c in cols)
 
     @abstractmethod
-    def apply_sql_clause(self, query):
+    def apply_sql_clause(self, query: Any) -> Any:
         pass
 
 
+@frozen
+class RegenerateSystemColumns(Step):
+    catalog: "Catalog"
+
+    def hash_inputs(self) -> str:
+        return hashlib.sha256(b"regenerate_system_columns").hexdigest()
+
+    def apply(
+        self, query_generator: QueryGenerator, temp_tables: list[str]
+    ) -> StepResult:
+        query = query_generator.select()
+        new_query = self.catalog.warehouse._regenerate_system_columns(
+            query, keep_existing_columns=True
+        )
+
+        def q(*columns):
+            return new_query.with_only_columns(*columns)
+
+        return step_result(q, new_query.selected_columns)
+
+
 @frozen
 class SQLSelect(SQLClause):
     args: tuple[Function | ColumnElement, ...]
@@ -1488,10 +1509,6 @@ class DatasetQuery:
         finally:
             self.cleanup()
 
-    def shuffle(self) -> "Self":
-        # ToDo: implement shaffle based on seed and/or generating random column
-        return self.order_by(C.sys__rand)
-
     def sample(self, n) -> "Self":
         """
         Return a random sample from the dataset.

{datachain-0.37.7 → datachain-0.37.8}/src/datachain/toolkit/split.py

@@ -1,6 +1,7 @@
 import random
 
 from datachain import C, DataChain
+from datachain.lib.signal_schema import SignalResolvingError
 
 RESOLUTION = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
 
@@ -59,7 +60,10 @@ def train_test_split(
         ```
 
     Note:
-        The splits are random but deterministic, based on Dataset `sys__rand` field.
+        Splits reuse the same best-effort shuffle used by `DataChain.shuffle`. Results
+        are typically repeatable, but earlier operations such as `merge`, `union`, or
+        custom SQL that reshuffle rows can change the outcome between runs. Add order by
+        stable keys first when you need strict reproducibility.
     """
     if len(weights) < 2:
         raise ValueError("Weights should have at least two elements")
@@ -68,16 +72,34 @@ def train_test_split(
 
     weights_normalized = [weight / sum(weights) for weight in weights]
 
+    try:
+        dc.signals_schema.resolve("sys.rand")
+    except SignalResolvingError:
+        dc = dc.persist()
+
     rand_col = C("sys.rand")
     if seed is not None:
         uniform_seed = random.Random(seed).randrange(1, RESOLUTION)  # noqa: S311
         rand_col = (rand_col % RESOLUTION) * uniform_seed  # type: ignore[assignment]
     rand_col = rand_col % RESOLUTION  # type: ignore[assignment]
 
-    return [
-        dc.filter(
-            rand_col >= round(sum(weights_normalized[:index]) * (RESOLUTION - 1)),
-            rand_col < round(sum(weights_normalized[: index + 1]) * (RESOLUTION - 1)),
-        )
-        for index, _ in enumerate(weights_normalized)
-    ]
+    boundaries: list[int] = [0]
+    cumulative = 0.0
+    for weight in weights_normalized[:-1]:
+        cumulative += weight
+        boundary = round(cumulative * RESOLUTION)
+        boundaries.append(min(boundary, RESOLUTION))
+    boundaries.append(RESOLUTION)
+
+    splits: list[DataChain] = []
+    last_index = len(weights_normalized) - 1
+    for index in range(len(weights_normalized)):
+        lower = boundaries[index]
+        if index == last_index:
+            condition = rand_col >= lower
+        else:
+            upper = boundaries[index + 1]
+            condition = (rand_col >= lower) & (rand_col < upper)
+        splits.append(dc.filter(condition))
+
+    return splits

{datachain-0.37.7 → datachain-0.37.8}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.37.7
+Version: 0.37.8
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.37.7 → datachain-0.37.8}/tests/func/test_toolkit.py

@@ -1,5 +1,6 @@
 import pytest
 
+import datachain as dc
 from datachain.toolkit import train_test_split
 
 
@@ -18,8 +19,8 @@ def test_train_test_split_not_random(not_random_ds, seed, weights, expected):
     res = train_test_split(not_random_ds, weights, seed=seed)
     assert len(res) == len(expected)
 
-    for i, dc in enumerate(res):
-        assert dc.to_values("sys.id") == expected[i]
+    for i, chain in enumerate(res):
+        assert chain.to_values("sys.id") == expected[i]
 
 
 @pytest.mark.parametrize(
@@ -40,8 +41,8 @@ def test_train_test_split_random(pseudo_random_ds, seed, weights, expected):
     res = train_test_split(pseudo_random_ds, weights, seed=seed)
     assert len(res) == len(expected)
 
-    for i, dc in enumerate(res):
-        assert dc.to_values("sys.id") == expected[i]
+    for i, chain in enumerate(res):
+        assert chain.to_values("sys.id") == expected[i]
 
 
 def test_train_test_split_errors(not_random_ds):
@@ -49,3 +50,32 @@ def test_train_test_split_errors(not_random_ds):
         train_test_split(not_random_ds, [0.5])
     with pytest.raises(ValueError, match="Weights should be non-negative"):
         train_test_split(not_random_ds, [-1, 1])
+
+
+def test_split_after_merge(test_session):
+    left = dc.read_values(ids=[1, 2, 3, 4], session=test_session)
+    right = dc.read_values(
+        ids=[1, 2, 3, 4],
+        extra=["a", "b", "c", "d"],
+        session=test_session,
+    )
+
+    merged = left.merge(right, on="ids")
+
+    train, test = train_test_split(merged, [0.5, 0.5])
+
+    for split in (train, test):
+        sys_schema = split.signals_schema.resolve("sys.id", "sys.rand").values
+        assert sys_schema["sys.id"] is int
+        assert sys_schema["sys.rand"] is int
+
+    combined_rows = set(train.to_list("ids", "extra")) | set(
+        test.to_list("ids", "extra")
+    )
+
+    assert combined_rows == {
+        (1, "a"),
+        (2, "b"),
+        (3, "c"),
+        (4, "d"),
+    }

{datachain-0.37.7 → datachain-0.37.8}/tests/unit/lib/test_datachain.py

@@ -1235,6 +1235,20 @@ def test_persist_restores_sys_signals_after_merge(test_session):
     assert sys_schema["sys.rand"] is int
 
 
+def test_shuffle_after_merge(test_session):
+    left = dc.read_values(ids=[1, 2], session=test_session)
+    right = dc.read_values(ids=[1, 2], extra=["x", "y"], session=test_session)
+
+    shuffled = left.merge(right, on="ids").shuffle()
+
+    sys_schema = shuffled.signals_schema.resolve("sys.id", "sys.rand").values
+    assert sys_schema["sys.id"] is int
+    assert sys_schema["sys.rand"] is int
+
+    rows = set(shuffled.to_list("ids", "extra"))
+    assert rows == {(1, "x"), (2, "y")}
+
+
 def test_unsupported_output_type(test_session):
     vector = [3.14, 2.72, 1.62]