datachain 0.37.4__tar.gz → 0.37.6__tar.gz

{datachain-0.37.4 → datachain-0.37.6}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.14.1'
+    rev: 'v0.14.2'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.37.4 → datachain-0.37.6}/PKG-INFO

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

{datachain-0.37.4 → datachain-0.37.6}/src/datachain/query/dataset.py

@@ -426,8 +426,10 @@ class UDFStep(Step, ABC):
         """Method that creates a table where temp udf results will be saved"""
 
     def process_input_query(self, query: Select) -> tuple[Select, list["Table"]]:
-        """Apply any necessary processing to the input query"""
-        return query, []
+        """Materialize inputs, ensure sys columns are available, needed for checkpoints,
+        needed for map to work (merge results)"""
+        table = self.catalog.warehouse.create_pre_udf_table(query)
+        return sqlalchemy.select(*table.c), [table]
 
     @abstractmethod
     def create_result_query(
@@ -628,23 +630,18 @@ class UDFStep(Step, ABC):
     def apply(
         self, query_generator: QueryGenerator, temp_tables: list[str]
     ) -> "StepResult":
-        _query = query = query_generator.select()
+        query, tables = self.process_input_query(query_generator.select())
+        _query = query
 
         # Apply partitioning if needed.
         if self.partition_by is not None:
-            _query = query = self.catalog.warehouse._regenerate_system_columns(
-                query_generator.select(),
-                keep_existing_columns=True,
-                regenerate_columns=["sys__id"],
-            )
             partition_tbl = self.create_partitions_table(query)
-            temp_tables.append(partition_tbl.name)
             query = query.outerjoin(
                 partition_tbl,
                 partition_tbl.c.sys__id == query.selected_columns.sys__id,
             ).add_columns(*partition_columns())
+            tables = [*tables, partition_tbl]
 
-        query, tables = self.process_input_query(query)
         temp_tables.extend(t.name for t in tables)
         udf_table = self.create_udf_table(_query)
         temp_tables.append(udf_table.name)
@@ -675,13 +672,6 @@ class UDFSignal(UDFStep):
 
         return self.catalog.warehouse.create_udf_table(udf_output_columns)
 
-    def process_input_query(self, query: Select) -> tuple[Select, list["Table"]]:
-        if os.getenv("DATACHAIN_DISABLE_QUERY_CACHE", "") not in ("", "0"):
-            return query, []
-        table = self.catalog.warehouse.create_pre_udf_table(query)
-        q: Select = sqlalchemy.select(*table.c)
-        return q, [table]
-
     def create_result_query(
         self, udf_table, query
     ) -> tuple[QueryGeneratorFunc, list["sqlalchemy.Column"]]:

{datachain-0.37.4 → datachain-0.37.6}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.37.4 → datachain-0.37.6}/tests/func/test_datachain.py

@@ -1699,10 +1699,16 @@ def test_agg_offset_limit(catalog_tmpfile, parallel, offset, limit, files):
         value=list(range(100)),
         session=catalog_tmpfile.session,
     )
+    # Read values in general doesn't guarantee order, so we need to order first
+    ds = ds.order_by("filename")
     if offset is not None:
         ds = ds.offset(offset)
     if limit is not None:
         ds = ds.limit(limit)
+
+    limited_filenames = ds.to_values("filename")
+    assert set(limited_filenames) == set(files)
+
     ds = (
         ds.settings(parallel=parallel)
         .agg(

{datachain-0.37.4 → datachain-0.37.6}/tests/func/test_udf.py

@@ -878,3 +878,146 @@ def test_udf_distributed_interrupt(
         chain.show()
     captured = capfd.readouterr()
     assert "semaphore" not in captured.err
+
+
+def test_gen_works_after_union(test_session_tmpfile, monkeypatch):
+    """
+    Union drops sys columns, we test that UDF generates them correctly after that.
+    """
+    monkeypatch.setattr("datachain.query.dispatch.DEFAULT_BATCH_SIZE", 5, raising=False)
+    n = 30
+
+    x_ids = list(range(n))
+    y_ids = list(range(n, 2 * n))
+
+    x = dc.read_values(idx=x_ids, session=test_session_tmpfile)
+    y = dc.read_values(idx=y_ids, session=test_session_tmpfile)
+
+    xy = x.union(y)
+
+    def expand(idx):
+        yield f"val-{idx}"
+
+    generated = xy.settings(parallel=2).gen(
+        gen=expand,
+        params=("idx",),
+        output={"val": str},
+    )
+
+    values = generated.to_values("val")
+
+    assert len(values) == 2 * n
+    assert set(values) == {f"val-{i}" for i in range(2 * n)}
+
+
+@pytest.mark.parametrize("full", [False, True])
+def test_gen_works_after_merge(test_session_tmpfile, monkeypatch, full):
+    """Merge drops sys columns as well; ensure UDF generation still works."""
+    monkeypatch.setattr("datachain.query.dispatch.DEFAULT_BATCH_SIZE", 5, raising=False)
+    n = 30
+
+    idxs = list(range(n))
+
+    left = dc.read_values(
+        idx=idxs,
+        left_value=[f"left-{i}" for i in idxs],
+        session=test_session_tmpfile,
+    )
+    right = dc.read_values(
+        idx=idxs,
+        right_value=[f"right-{i}" for i in idxs],
+        session=test_session_tmpfile,
+    )
+
+    merged = left.merge(right, on="idx", full=full)
+
+    def expand(idx, left_value, right_value):
+        yield f"val-{idx}-{left_value}-{right_value}"
+
+    generated = merged.settings(parallel=2).gen(
+        gen=expand,
+        params=("idx", "left_value", "right_value"),
+        output={"val": str},
+    )
+
+    values = generated.to_values("val")
+
+    assert len(values) == n
+    expected = {f"val-{i}-left-{i}-right-{i}" for i in idxs}
+    assert set(values) == expected
+
+
+def test_agg_works_after_union(test_session_tmpfile, monkeypatch):
+    """Union must preserve sys columns for aggregations with functional partitions."""
+    from datachain import func
+
+    monkeypatch.setattr("datachain.query.dispatch.DEFAULT_BATCH_SIZE", 5, raising=False)
+
+    groups = 5
+    n = 30
+
+    x_paths = [f"group-{i % groups}/item-{i}" for i in range(n)]
+    y_paths = [f"group-{i % groups}/item-{n + i}" for i in range(n)]
+
+    x = dc.read_values(path=x_paths, session=test_session_tmpfile)
+    y = dc.read_values(path=y_paths, session=test_session_tmpfile)
+
+    xy = x.union(y)
+
+    def summarize(paths):
+        group = paths[0].split("/")[0]
+        yield group, len(paths)
+
+    aggregated = xy.settings(parallel=2).agg(
+        summarize,
+        params=("path",),
+        output={"partition": str, "count": int},
+        partition_by=func.parent("path"),
+    )
+
+    records = aggregated.to_records()
+    expected_counts = {f"group-{g}": 2 * n // groups for g in range(groups)}
+    assert {row["partition"]: row["count"] for row in records} == expected_counts
+
+
+@pytest.mark.parametrize("full", [False, True])
+def test_agg_works_after_merge(test_session_tmpfile, monkeypatch, full):
+    """Ensure merge keeps sys columns for aggregations with functional partitions."""
+    from datachain import func
+
+    monkeypatch.setattr("datachain.query.dispatch.DEFAULT_BATCH_SIZE", 5, raising=False)
+
+    groups = 5
+    n = 30
+    idxs = list(range(n))
+
+    left = dc.read_values(
+        idx=idxs,
+        left_path=[f"group-{i % groups}/left-{i}" for i in idxs],
+        session=test_session_tmpfile,
+    )
+    right = dc.read_values(
+        idx=idxs,
+        right_value=idxs,
+        session=test_session_tmpfile,
+    )
+
+    merged = left.merge(right, on="idx", full=full)
+
+    def summarize(left_path, right_value):
+        group = left_path[0].split("/")[0]
+        yield group, sum(right_value)
+
+    aggregated = merged.settings(parallel=2).agg(
+        summarize,
+        params=("left_path", "right_value"),
+        output={"partition": str, "total": int},
+        partition_by=func.parent("left_path"),
+    )
+
+    records = aggregated.to_records()
+    expected_totals = {
+        f"group-{g}": sum(val for val in idxs if val % groups == g)
+        for g in range(groups)
+    }
+    assert {row["partition"]: row["total"] for row in records} == expected_totals