Making listing lazy in DatasetQuery (#976)

* adding listing as pre-step

* Update src/datachain/lib/dc.py

Co-authored-by: Vladimir Rudnykh <[email protected]>

* Update src/datachain/query/dataset.py

Co-authored-by: Vladimir Rudnykh <[email protected]>

* returned to starting step

---------

Co-authored-by: Vladimir Rudnykh <[email protected]>

Author: ilongin

src/datachain/catalog/catalog.py

@@ -588,7 +588,7 @@ def enlist_source(
 
         from_storage(
             source, session=self.session, update=update, object_name=object_name
-        )
+        ).exec()
 
         list_ds_name, list_uri, list_path, _ = get_listing(
             source, self.session, update=update

src/datachain/lib/dc/storage.py

@@ -6,7 +6,6 @@
 )
 
 from datachain.lib.file import (
-    File,
     FileType,
     get_file_type,
 )
@@ -95,24 +94,28 @@ def from_storage(
         dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
         return dc
 
+    dc = from_dataset(list_ds_name, session=session, settings=settings)
+    dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
+
     if update or not list_ds_exists:
-        # disable prefetch for listing, as it pre-downloads all files
-        (
-            from_records(
-                DataChain.DEFAULT_FILE_RECORD,
-                session=session,
-                settings=settings,
-                in_memory=in_memory,
-            )
-            .settings(prefetch=0)
-            .gen(
-                list_bucket(list_uri, cache, client_config=client_config),
-                output={f"{object_name}": File},
+
+        def lst_fn():
+            # disable prefetch for listing, as it pre-downloads all files
+            (
+                from_records(
+                    DataChain.DEFAULT_FILE_RECORD,
+                    session=session,
+                    settings=settings,
+                    in_memory=in_memory,
+                )
+                .settings(prefetch=0)
+                .gen(
+                    list_bucket(list_uri, cache, client_config=client_config),
+                    output={f"{object_name}": file_type},
+                )
+                .save(list_ds_name, listing=True)
             )
-            .save(list_ds_name, listing=True)
-        )
 
-    dc = from_dataset(list_ds_name, session=session, settings=settings)
-    dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
+        dc._query.add_before_steps(lst_fn)
 
     return ls(dc, list_path, recursive=recursive, object_name=object_name)

src/datachain/query/dataset.py

@@ -47,6 +47,7 @@
     QueryScriptCancelError,
 )
 from datachain.func.base import Function
+from datachain.lib.listing import is_listing_dataset
 from datachain.lib.udf import UDFAdapter, _get_cache
 from datachain.progress import CombinedDownloadCallback, TqdmCombinedDownloadCallback
 from datachain.query.schema import C, UDFParamSpec, normalize_param
@@ -151,13 +152,6 @@ def step_result(
     )
 
 
-class StartingStep(ABC):
-    """An initial query processing step, referencing a data source."""
-
-    @abstractmethod
-    def apply(self) -> "StepResult": ...
-
-
 @frozen
 class Step(ABC):
     """A query processing step (filtering, mutation, etc.)"""
@@ -170,7 +164,7 @@ def apply(
 
 
 @frozen
-class QueryStep(StartingStep):
+class QueryStep:
     catalog: "Catalog"
     dataset_name: str
     dataset_version: int
@@ -1097,26 +1091,42 @@ def __init__(
         self.temp_table_names: list[str] = []
         self.dependencies: set[DatasetDependencyType] = set()
         self.table = self.get_table()
-        self.starting_step: StartingStep
+        self.starting_step: Optional[QueryStep] = None
         self.name: Optional[str] = None
         self.version: Optional[int] = None
         self.feature_schema: Optional[dict] = None
         self.column_types: Optional[dict[str, Any]] = None
+        self.before_steps: list[Callable] = []
 
-        self.name = name
+        self.list_ds_name: Optional[str] = None
 
-        if fallback_to_studio and is_token_set():
-            ds = self.catalog.get_dataset_with_remote_fallback(name, version)
+        self.name = name
+        self.dialect = self.catalog.warehouse.db.dialect
+        if version:
+            self.version = version
+
+        if is_listing_dataset(name):
+            # not setting query step yet as listing dataset might not exist at
+            # this point
+            self.list_ds_name = name
+        elif fallback_to_studio and is_token_set():
+            self._set_starting_step(
+                self.catalog.get_dataset_with_remote_fallback(name, version)
+            )
         else:
-            ds = self.catalog.get_dataset(name)
+            self._set_starting_step(self.catalog.get_dataset(name))
+
+    def _set_starting_step(self, ds: "DatasetRecord") -> None:
+        if not self.version:
+            self.version = ds.latest_version
 
-        self.version = version or ds.latest_version
+        self.starting_step = QueryStep(self.catalog, ds.name, self.version)
+
+        # at this point we know our starting dataset so setting up schemas
         self.feature_schema = ds.get_version(self.version).feature_schema
         self.column_types = copy(ds.schema)
         if "sys__id" in self.column_types:
             self.column_types.pop("sys__id")
-        self.starting_step = QueryStep(self.catalog, name, self.version)
-        self.dialect = self.catalog.warehouse.db.dialect
 
     def __iter__(self):
         return iter(self.db_results())
@@ -1180,11 +1190,23 @@ def c(self, column: Union[C, str]) -> "ColumnClause[Any]":
         col.table = self.table
         return col
 
+    def add_before_steps(self, fn: Callable) -> None:
+        """
+        Setting custom function to be run before applying steps
+        """
+        self.before_steps.append(fn)
+
     def apply_steps(self) -> QueryGenerator:
         """
         Apply the steps in the query and return the resulting
         sqlalchemy.SelectBase.
         """
+        for fn in self.before_steps:
+            fn()
+
+        if self.list_ds_name:
+            # at this point we know what is our starting listing dataset name
+            self._set_starting_step(self.catalog.get_dataset(self.list_ds_name))  # type: ignore [arg-type]
         query = self.clone()
 
         index = os.getenv("DATACHAIN_QUERY_CHUNK_INDEX", self._chunk_index)
@@ -1203,6 +1225,7 @@ def apply_steps(self) -> QueryGenerator:
             query = query.filter(C.sys__rand % total == index)
             query.steps = query.steps[-1:] + query.steps[:-1]
 
+        assert query.starting_step
         result = query.starting_step.apply()
         self.dependencies.update(result.dependencies)
 

tests/func/test_datachain.py

@@ -152,7 +152,7 @@ def _list_dataset_name(uri: str) -> str:
         return name
 
     dogs_uri = f"{src_uri}/dogs"
-    dc.from_storage(dogs_uri, session=session)
+    dc.from_storage(dogs_uri, session=session).exec()
     assert _get_listing_datasets(session) == [
         f"{_list_dataset_name(dogs_uri)}@v1",
     ]
@@ -162,15 +162,15 @@ def _list_dataset_name(uri: str) -> str:
         f"{_list_dataset_name(dogs_uri)}@v1",
     ]
 
-    dc.from_storage(src_uri, session=session)
+    dc.from_storage(src_uri, session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(dogs_uri)}@v1",
             f"{_list_dataset_name(src_uri)}@v1",
         ]
     )
 
-    dc.from_storage(f"{src_uri}/cats", session=session)
+    dc.from_storage(f"{src_uri}/cats", session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(dogs_uri)}@v1",
@@ -196,14 +196,14 @@ def _list_dataset_name(uri: str) -> str:
         return name
 
     uri = f"{src_uri}/cats"
-    dc.from_storage(uri, session=session)
+    dc.from_storage(uri, session=session).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(uri)}@v1",
         ]
     )
 
-    dc.from_storage(uri, session=session, update=True)
+    dc.from_storage(uri, session=session, update=True).exec()
     assert _get_listing_datasets(session) == sorted(
         [
             f"{_list_dataset_name(uri)}@v1",

tests/func/test_ls.py

@@ -32,7 +32,7 @@ def test_ls_no_args(cloud_test_catalog, cloud_type, capsys):
     catalog = session.catalog
     src = cloud_test_catalog.src_uri
 
-    dc.from_storage(src, session=session).collect()
+    dc.from_storage(src, session=session).exec()
     ls([], catalog=catalog)
     captured = capsys.readouterr()
     assert captured.out == f"{src}/@v1\n"

tests/unit/lib/test_datachain.py

@@ -339,7 +339,7 @@ def test_listings(test_session, tmp_dir):
     df.to_parquet(tmp_dir / "df.parquet")
 
     uri = tmp_dir.as_uri()
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
 
     # check that listing is not returned as normal dataset
     assert not any(
@@ -370,13 +370,13 @@ def test_listings_reindex(test_session, tmp_dir):
 
     uri = tmp_dir.as_uri()
 
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
     assert len(list(dc.listings(session=test_session).collect("listing"))) == 1
 
-    dc.from_storage(uri, session=test_session)
+    dc.from_storage(uri, session=test_session).exec()
     assert len(list(dc.listings(session=test_session).collect("listing"))) == 1
 
-    dc.from_storage(uri, session=test_session, update=True)
+    dc.from_storage(uri, session=test_session, update=True).exec()
     listings = list(dc.listings(session=test_session).collect("listing"))
     assert len(listings) == 2
     listings.sort(key=lambda lst: lst.version)