datachain 0.1.11__tar.gz → 0.1.12__tar.gz

{datachain-0.1.11/src/datachain.egg-info → datachain-0.1.12}/PKG-INFO

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

{datachain-0.1.11 → datachain-0.1.12}/src/datachain/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.11'
-__version_tuple__ = version_tuple = (0, 1, 11)
+__version__ = version = '0.1.12'
+__version_tuple__ = version_tuple = (0, 1, 12)

{datachain-0.1.11 → datachain-0.1.12}/src/datachain/lib/dc.py

@@ -492,16 +492,43 @@ class DataChain(DatasetQuery):
         chain.signals_schema = new_schema
         return chain
 
-    def get_values(self) -> Iterator[list]:
-        """Iterate over rows, getting feature values and applying reader calls."""
-        for features in self.iterate():
-            yield [fr.get_value() if isinstance(fr, Feature) else fr for fr in features]
+    def get_values(self, *cols: str) -> Iterator[list]:
+        """Iterate over rows, getting feature values and applying reader calls.
+        If columns are specified - limit them to specified columns.
+        """
+        for features in self.iterate(*cols):
+            yield [fr.get_value() if isinstance(fr, Feature) else fr for fr in features]  # type: ignore[union-attr,call-arg]
+
+    def get_one_value(self, col: str) -> Iterator:
+        for item in self.get_values(col):
+            yield item[0]
 
-    def iterate(self) -> Iterator[Sequence[Feature]]:
-        db_signals = self.signals_schema.db_signals()
+    def iterate(self, *cols: str) -> Iterator[list[FeatureType]]:
+        """Iterate over rows. If columns are specified - limit them to specified
+        columns.
+        """
+        chain = self.select(*cols) if cols else self
+
+        db_signals = chain.signals_schema.db_signals()
         with super().select(*db_signals).as_iterable() as rows_iter:
             for row in rows_iter:
-                yield self.signals_schema.row_to_features(row, self.session.catalog)
+                yield chain.signals_schema.row_to_features(row, chain.session.catalog)
+
+    def iterate_one(self, col: str) -> Iterator[FeatureType]:
+        for item in self.iterate(col):
+            yield item[0]
+
+    def collect(self, *cols: str) -> list[list[FeatureType]]:
+        return list(self.iterate(*cols))
+
+    def collect_one(self, col: str) -> list[FeatureType]:
+        return list(self.iterate_one(col))
+
+    def collect_values(self, *cols: str) -> list[list]:
+        return list(self.get_values(*cols))
+
+    def collect_one_value(self, col: str) -> list:
+        return list(self.get_one_value(col))
 
     def to_pytorch(self, **kwargs):
         """Convert to pytorch dataset format."""

{datachain-0.1.11 → datachain-0.1.12/src/datachain.egg-info}/PKG-INFO

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

{datachain-0.1.11 → datachain-0.1.12}/tests/func/test_catalog.py

@@ -1078,6 +1078,7 @@ def test_garbage_collect(cloud_test_catalog, from_cli, capsys):
 
 
 def test_get_file_signals(cloud_test_catalog, dogs_dataset):
+    skip_if_not_sqlite()
     catalog = cloud_test_catalog.catalog
     catalog.metastore.update_dataset_version(
         dogs_dataset,
@@ -1105,6 +1106,7 @@ def test_get_file_signals(cloud_test_catalog, dogs_dataset):
 
 
 def test_get_file_signals_no_signals(cloud_test_catalog, dogs_dataset):
+    skip_if_not_sqlite()
     catalog = cloud_test_catalog.catalog
     catalog.metastore.update_dataset_version(
         dogs_dataset,
@@ -1123,6 +1125,7 @@ def test_get_file_signals_no_signals(cloud_test_catalog, dogs_dataset):
 
 
 def test_open_object_no_file_signals(cloud_test_catalog, dogs_dataset):
+    skip_if_not_sqlite()
     catalog = cloud_test_catalog.catalog
     catalog.metastore.update_dataset_version(
         dogs_dataset,

{datachain-0.1.11 → datachain-0.1.12}/tests/unit/lib/test_datachain.py

@@ -17,6 +17,7 @@ from datachain.lib.signal_schema import (
 )
 from datachain.lib.udf_signature import UdfSignatureError
 from datachain.lib.utils import DataChainParamsError
+from tests.utils import assert_vector_is_close
 
 DF_DATA = {
     "first_name": ["Alice", "Bob", "Charlie", "David", "Eva"],
@@ -212,8 +213,7 @@ def test_gen(catalog):
     df = ds.to_pandas()
 
     assert df["x__my_name"].tolist() == ["n1", "n2", "n1"]
-    for actual_sqrt, expected in zip(df["x__sqrt"], [3, 5, 1]):
-        assert math.isclose(actual_sqrt, math.sqrt(expected), rel_tol=1e-7)
+    assert_vector_is_close(df["x__sqrt"], [3, 5, 1], expected_callback=math.sqrt)
     with pytest.raises(KeyError):
         df["x__t1__nnn"]
 
@@ -235,8 +235,7 @@ def test_map(catalog):
     ).to_pandas()
 
     assert df["x__my_name"].tolist() == ["n1_suf", "n2_suf", "n1_suf"]
-    for actual_sqrt, expected in zip(df["x__sqrt"], [3, 5, 1]):
-        assert math.isclose(actual_sqrt, math.sqrt(expected), rel_tol=1e-7)
+    assert_vector_is_close(df["x__sqrt"], [3, 5, 1], expected_callback=math.sqrt)
 
 
 def test_agg(catalog):
@@ -384,7 +383,7 @@ def test_agg_tuple_result_generator(catalog):
     assert df["x_1__size"].tolist() == [10, 5]
 
 
-def test_iterate():
+def test_iterate(catalog):
     dc = DataChain.from_features(f1=features, num=range(len(features)))
 
     n = 0
@@ -402,7 +401,7 @@ def test_iterate():
     assert n == len(features)
 
 
-def test_iterate_nested_feature():
+def test_iterate_nested_feature(catalog):
     dc = DataChain.from_features(sign1=features_nested)
 
     for n, sample in enumerate(dc.iterate()):
@@ -413,7 +412,7 @@ def test_iterate_nested_feature():
         assert nested == features_nested[n]
 
 
-def test_select_feature():
+def test_select_feature(catalog):
     dc = DataChain.from_features(my_n=features_nested)
 
     samples = dc.select("my_n").iterate()
@@ -440,7 +439,7 @@ def test_select_feature():
     assert n == len(features_nested)
 
 
-def test_select_columns_intersection():
+def test_select_columns_intersection(catalog):
     dc = DataChain.from_features(my_n=features_nested)
 
     samples = dc.select("my_n.fr", "my_n.fr.count").iterate()
@@ -453,7 +452,7 @@ def test_select_columns_intersection():
     assert n == len(features_nested)
 
 
-def test_select_except():
+def test_select_except(catalog):
     dc = DataChain.from_features(fr1=features_nested, fr2=features)
 
     samples = dc.select_except("fr2").iterate()
@@ -465,7 +464,7 @@ def test_select_except():
     assert n == len(features_nested)
 
 
-def test_select_wrong_type():
+def test_select_wrong_type(catalog):
     dc = DataChain.from_features(fr1=features_nested, fr2=features)
 
     with pytest.raises(SignalResolvingTypeError):
@@ -475,7 +474,7 @@ def test_select_wrong_type():
         list(dc.select_except(features[0]).iterate())
 
 
-def test_select_except_error():
+def test_select_except_error(catalog):
     dc = DataChain.from_features(fr1=features_nested, fr2=features)
 
     with pytest.raises(SignalResolvingError):
@@ -485,7 +484,7 @@ def test_select_except_error():
         list(dc.select_except("fr1.label", "file").iterate())
 
 
-def test_select_restore_from_saving():
+def test_select_restore_from_saving(catalog):
     dc = DataChain.from_features(my_n=features_nested)
 
     name = "test_test_select_save"
@@ -493,13 +492,15 @@ def test_select_restore_from_saving():
 
     restored = DataChain.from_dataset(name)
     n = 0
-    for sample in restored.iterate():
-        assert sample[0] == features[n]
+    restored_sorted = sorted(restored.iterate(), key=lambda x: x[0].count)
+    features_sorted = sorted(features, key=lambda x: x.count)
+    for sample in restored_sorted:
+        assert sample[0] == features_sorted[n]
         n += 1
     assert n == len(features_nested)
 
 
-def test_chain_of_maps():
+def test_chain_of_maps(catalog):
     dc = (
         DataChain.from_features(my_n=features_nested)
         .map(full_name=lambda my_n: my_n.label + "-" + my_n.fr.nnn, output=str)
@@ -516,7 +517,7 @@ def test_chain_of_maps():
         assert signal in preserved.signals_schema.values
 
 
-def test_vector():
+def test_vector(catalog):
     vector = [3.14, 2.72, 1.62]
 
     def get_vector(key) -> list[float]:
@@ -525,10 +526,10 @@ def test_vector():
     ds = DataChain.from_features(key=[123]).map(emd=get_vector)
 
     df = ds.to_pandas()
-    assert df["emd"].tolist()[0] == vector
+    assert_vector_is_close(df["emd"].tolist()[0], vector)
 
 
-def test_vector_of_vectors():
+def test_vector_of_vectors(catalog):
     vector = [[3.14, 2.72, 1.62], [1.0, 2.0, 3.0]]
 
     def get_vector(key) -> list[list[float]]:
@@ -537,10 +538,13 @@ def test_vector_of_vectors():
     ds = DataChain.from_features(key=[123]).map(emd_list=get_vector)
 
     df = ds.to_pandas()
-    assert df["emd_list"].tolist()[0] == vector
+    actual = df["emd_list"].tolist()[0]
+    assert len(actual) == 2
+    assert_vector_is_close(actual[0], vector[0])
+    assert_vector_is_close(actual[1], vector[1])
 
 
-def test_unsupported_output_type():
+def test_unsupported_output_type(catalog):
     vector = [3.14, 2.72, 1.62]
 
     def get_vector(key) -> list[np.float64]:
@@ -548,3 +552,45 @@ def test_unsupported_output_type():
 
     with pytest.raises(SignalSchemaError):
         DataChain.from_features(key=[123]).map(emd=get_vector)
+
+
+def test_collect_one(catalog):
+    names = ["f1.jpg", "f1.json", "f1.txt", "f2.jpg", "f2.json"]
+    sizes = [1, 2, 3, 4, 5]
+    files = [File(name=name, size=size) for name, size in zip(names, sizes)]
+
+    scores = [0.1, 0.2, 0.3, 0.4, 0.5]
+
+    chain = DataChain.from_features(file=files, score=scores)
+
+    assert chain.collect_one("file") == files
+    assert chain.collect_one("file.name") == names
+    assert chain.collect_one("file.size") == sizes
+    assert chain.collect_one("file.source") == [""] * len(names)
+    assert_vector_is_close(chain.collect_one("score"), scores)
+
+    for actual, expected in zip(
+        chain.collect("file.size", "score"), [[x, y] for x, y in zip(sizes, scores)]
+    ):
+        assert len(actual) == 2
+        assert actual[0] == expected[0]
+        assert math.isclose(actual[1], expected[1], rel_tol=1e-7)
+
+
+def test_get_values(catalog):
+    class _MyFr(Feature):
+        num: int
+
+        def get_value(self):
+            return self.num * 2
+
+    lst1 = [_MyFr(num=i) for i in range(6)]
+    lst2 = [_MyFr(num=-i) for i in range(6)]
+    chain = DataChain.from_features(fr1=lst1, fr2=lst2)
+
+    assert list(chain.get_one_value("fr1")) == [i * 2 for i in range(6)]
+    assert chain.collect_one_value("fr1") == [i * 2 for i in range(6)]
+
+    assert chain.collect_values("fr1", "fr2") == [
+        [x.get_value(), y.get_value()] for x, y in zip(lst1, lst2)
+    ]

{datachain-0.1.11 → datachain-0.1.12}/tests/unit/lib/test_datachain_merge.py

@@ -1,3 +1,4 @@
+import math
 from typing import Optional
 
 import pytest
@@ -5,6 +6,7 @@ import pytest
 from datachain.lib.dc import DataChain, DatasetMergeError
 from datachain.lib.feature import Feature
 from datachain.lib.signal_schema import SignalResolvingError
+from datachain.sql.types import Float, String
 
 
 class TestUser(Feature):
@@ -42,11 +44,14 @@ team = [
 ]
 
 
-def test_merge_objects():
+def test_merge_objects(catalog):
     ch1 = DataChain.from_features(emp=employees)
     ch2 = DataChain.from_features(team=team)
     ch = ch1.merge(ch2, "emp.person.name", "team.player")
 
+    str_default = String.default_value(catalog.warehouse.db.dialect)
+    float_default = Float.default_value(catalog.warehouse.db.dialect)
+
     i = 0
     j = 0
     for items in ch.iterate():
@@ -59,16 +64,22 @@ def test_merge_objects():
 
         assert isinstance(player, TestTeamMember)
         if empl.person.name != "Bob":
-            assert player == team[j]
+            assert player.player == team[j].player
+            assert player.sport == team[j].sport
+            assert math.isclose(player.weight, team[j].weight, rel_tol=1e-7)
+            assert math.isclose(player.height, team[j].height, rel_tol=1e-7)
             j += 1
         else:
-            assert player == TestTeamMember()
+            assert player.player == str_default
+            assert player.sport == str_default
+            assert player.weight == float_default
+            assert player.height == float_default
 
     assert i == len(employees)
     assert j == len(team)
 
 
-def test_merge_similar_objects():
+def test_merge_similar_objects(catalog):
     new_employees = [
         TestEmployee(id=152, person=TestUser(name="Bob", age=27)),
         TestEmployee(id=201, person=TestUser(name="Karl", age=18)),
@@ -91,13 +102,15 @@ def test_merge_similar_objects():
     assert len(list(ch_inner.iterate())) == 2
 
 
-def test_merge_values():
+def test_merge_values(catalog):
     order_ids = [11, 22, 33, 44]
     order_descr = ["water", "water", "paper", "water"]
 
     delivery_ids = [11, 44]
     delivery_time = [24.0, 16.5]
 
+    float_default = Float.default_value(catalog.warehouse.db.dialect)
+
     ch1 = DataChain.from_features(id=order_ids, descr=order_descr)
     ch2 = DataChain.from_features(id=delivery_ids, time=delivery_time)
 
@@ -107,7 +120,8 @@ def test_merge_values():
 
     i = 0
     j = 0
-    for items in ch.iterate():
+    sorted_items_list = sorted(ch.iterate(), key=lambda x: x[0])
+    for items in sorted_items_list:
         assert len(items) == 4
         id, name, right_id, time = items
 
@@ -115,7 +129,7 @@ def test_merge_values():
         assert name == order_descr[i]
         i += 1
 
-        if time is not None:
+        if time != float_default:
             assert id == delivery_ids[j]
             assert time == delivery_time[j]
             j += 1
@@ -124,7 +138,7 @@ def test_merge_values():
     assert j == len(delivery_ids)
 
 
-def test_merge_multi_conditions():
+def test_merge_multi_conditions(catalog):
     order_ids = [11, 22, 33, 44]
     order_name = ["water", "water", "paper", "water"]
     order_descr = ["still water", "still water", "white paper", "sparkling water"]
@@ -151,7 +165,7 @@ def test_merge_multi_conditions():
     assert success_ids == {11}
 
 
-def test_merge_errors():
+def test_merge_errors(catalog):
     ch1 = DataChain.from_features(emp=employees)
     ch2 = DataChain.from_features(team=team)
 

{datachain-0.1.11 → datachain-0.1.12}/tests/unit/lib/test_feature_utils.py

@@ -21,7 +21,7 @@ def test_basic():
     assert vals[-1] == (fib[-1], values[-1])
 
 
-def test_e2e():
+def test_e2e(catalog):
     fib = [1, 1, 2, 3, 5, 8]
     values = ["odd" if num % 2 else "even" for num in fib]
 
@@ -44,7 +44,7 @@ def test_single_value():
     assert vals == fib
 
 
-def test_single_e2e():
+def test_single_e2e(catalog):
     fib = [1, 1, 2, 3, 5, 8]
 
     dc = DataChain.from_features(fib=fib)

{datachain-0.1.11 → datachain-0.1.12}/tests/utils.py

@@ -3,6 +3,7 @@ import io
 import math
 import os
 import tarfile
+from collections.abc import Iterable
 from string import printable
 from tarfile import DIRTYPE, TarInfo
 from time import sleep, time
@@ -282,3 +283,17 @@ def assert_row_names(
         == {r.get("name") for r in preview}
         == expected_names
     )
+
+
+def assert_vector_is_close(
+    actual: Iterable,
+    expected: Iterable,
+    expected_callback: Optional[Callable] = None,
+    rel_tol: float = 1e-7,
+) -> None:
+    """Asserts that two vectors of floating point numbers are equivalent, by ensuring
+    that all values are close together (within floating point tolerances)."""
+    for act, expect in zip(actual, expected):
+        if expected_callback:
+            expect = expected_callback(expect)
+        assert math.isclose(act, expect, rel_tol=rel_tol)