Listing refactoring and optimizations (#1114)

Author: dreadatour

src/datachain/catalog/catalog.py

@@ -66,6 +66,7 @@
     )
     from datachain.dataset import DatasetListVersion
     from datachain.job import Job
+    from datachain.lib.listing_info import ListingInfo
     from datachain.listing import Listing
 
 logger = logging.getLogger("datachain")
@@ -1116,13 +1117,16 @@ def get_dataset_dependencies(
         return direct_dependencies
 
     def ls_datasets(
-        self, include_listing: bool = False, studio: bool = False
+        self,
+        prefix: Optional[str] = None,
+        include_listing: bool = False,
+        studio: bool = False,
     ) -> Iterator[DatasetListRecord]:
         from datachain.remote.studio import StudioClient
 
         if studio:
             client = StudioClient()
-            response = client.ls_datasets()
+            response = client.ls_datasets(prefix=prefix)
             if not response.ok:
                 raise DataChainError(response.message)
             if not response.data:
@@ -1133,6 +1137,8 @@ def ls_datasets(
                 for d in response.data
                 if not d.get("name", "").startswith(QUERY_DATASET_PREFIX)
             )
+        elif prefix:
+            datasets = self.metastore.list_datasets_by_prefix(prefix)
         else:
             datasets = self.metastore.list_datasets()
 
@@ -1142,39 +1148,55 @@ def ls_datasets(
 
     def list_datasets_versions(
         self,
+        prefix: Optional[str] = None,
         include_listing: bool = False,
+        with_job: bool = True,
         studio: bool = False,
     ) -> Iterator[tuple[DatasetListRecord, "DatasetListVersion", Optional["Job"]]]:
         """Iterate over all dataset versions with related jobs."""
         datasets = list(
-            self.ls_datasets(include_listing=include_listing, studio=studio)
+            self.ls_datasets(
+                prefix=prefix, include_listing=include_listing, studio=studio
+            )
         )
 
         # preselect dataset versions jobs from db to avoid multiple queries
-        jobs_ids: set[str] = {
-            v.job_id for ds in datasets for v in ds.versions if v.job_id
-        }
         jobs: dict[str, Job] = {}
-        if jobs_ids:
-            jobs = {j.id: j for j in self.metastore.list_jobs_by_ids(list(jobs_ids))}
+        if with_job:
+            jobs_ids: set[str] = {
+                v.job_id for ds in datasets for v in ds.versions if v.job_id
+            }
+            if jobs_ids:
+                jobs = {
+                    j.id: j for j in self.metastore.list_jobs_by_ids(list(jobs_ids))
+                }
 
         for d in datasets:
             yield from (
-                (d, v, jobs.get(str(v.job_id)) if v.job_id else None)
+                (d, v, jobs.get(str(v.job_id)) if with_job and v.job_id else None)
                 for v in d.versions
             )
 
-    def listings(self):
+    def listings(self, prefix: Optional[str] = None) -> list["ListingInfo"]:
         """
         Returns list of ListingInfo objects which are representing specific
         storage listing datasets
         """
-        from datachain.lib.listing import is_listing_dataset
+        from datachain.lib.listing import LISTING_PREFIX, is_listing_dataset
         from datachain.lib.listing_info import ListingInfo
 
+        if prefix and not prefix.startswith(LISTING_PREFIX):
+            prefix = LISTING_PREFIX + prefix
+
+        listing_datasets_versions = self.list_datasets_versions(
+            prefix=prefix,
+            include_listing=True,
+            with_job=False,
+        )
+
         return [
             ListingInfo.from_models(d, v, j)
-            for d, v, j in self.list_datasets_versions(include_listing=True)
+            for d, v, j in listing_datasets_versions
             if is_listing_dataset(d.name)
         ]
 

src/datachain/dataset.py

@@ -93,7 +93,7 @@ def dataset_name(self) -> str:
         if self.type == DatasetDependencyType.DATASET:
             return self.name
 
-        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"), {})
+        list_dataset_name, _, _ = parse_listing_uri(self.name.strip("/"))
         assert list_dataset_name
         return list_dataset_name
 

src/datachain/lib/listing.py

@@ -107,11 +107,10 @@ def _file_c(name: str) -> Column:
     return dc.filter(pathfunc.parent(_file_c("path")) == path.lstrip("/").rstrip("/*"))
 
 
-def parse_listing_uri(uri: str, client_config) -> tuple[str, str, str]:
+def parse_listing_uri(uri: str) -> tuple[str, str, str]:
     """
     Parsing uri and returns listing dataset name, listing uri and listing path
     """
-    client_config = client_config or {}
     storage_uri, path = Client.parse_url(uri)
     if uses_glob(path):
         lst_uri_path = posixpath.dirname(path)
@@ -175,7 +174,7 @@ def get_listing(
         _, path = Client.parse_url(uri)
         return None, uri, path, False
 
-    ds_name, list_uri, list_path = parse_listing_uri(uri, client_config)
+    ds_name, list_uri, list_path = parse_listing_uri(uri)
     listing = None
     listings = [
         ls for ls in catalog.listings() if not ls.is_expired and ls.contains(ds_name)

src/datachain/remote/studio.py

@@ -282,8 +282,10 @@ def ls(self, paths: Iterable[str]) -> Iterator[tuple[str, Response[LsData]]]:
             response = self._send_request_msgpack("datachain/ls", {"source": path})
             yield path, response
 
-    def ls_datasets(self) -> Response[LsData]:
-        return self._send_request("datachain/datasets", {}, method="GET")
+    def ls_datasets(self, prefix: Optional[str] = None) -> Response[LsData]:
+        return self._send_request(
+            "datachain/datasets", {"prefix": prefix}, method="GET"
+        )
 
     def edit_dataset(
         self,

tests/func/test_catalog.py

@@ -15,7 +15,7 @@
 
 
 def listing_stats(uri, catalog):
-    list_dataset_name, _, _ = parse_listing_uri(uri, catalog.client_config)
+    list_dataset_name, _, _ = parse_listing_uri(uri)
     dataset = catalog.get_dataset(list_dataset_name)
     dataset_version = dataset.get_version(dataset.latest_version)
     return dataset_version.num_objects, dataset_version.size

tests/func/test_datachain.py

@@ -114,12 +114,11 @@ def test_read_storage_reindex(tmp_dir, test_session):
 
 
 def test_read_storage_reindex_expired(tmp_dir, test_session):
-    catalog = test_session.catalog
     tmp_dir = tmp_dir / "parquets"
     os.mkdir(tmp_dir)
     uri = tmp_dir.as_uri()
 
-    lst_ds_name = parse_listing_uri(uri, catalog.client_config)[0]
+    lst_ds_name = parse_listing_uri(uri)[0]
 
     pd.DataFrame({"name": ["Alice", "Bob"]}).to_parquet(tmp_dir / "test1.parquet")
     assert dc.read_storage(uri, session=test_session).count() == 1
@@ -144,10 +143,9 @@ def test_read_storage_partials(cloud_test_catalog):
     ctc = cloud_test_catalog
     src_uri = ctc.src_uri
     session = ctc.session
-    catalog = session.catalog
 
     def _list_dataset_name(uri: str) -> str:
-        name = parse_listing_uri(uri, catalog.client_config)[0]
+        name = parse_listing_uri(uri)[0]
         assert name
         return name
 
@@ -188,10 +186,9 @@ def test_read_storage_partials_with_update(cloud_test_catalog):
     ctc = cloud_test_catalog
     src_uri = ctc.src_uri
     session = ctc.session
-    catalog = session.catalog
 
     def _list_dataset_name(uri: str) -> str:
-        name = parse_listing_uri(uri, catalog.client_config)[0]
+        name = parse_listing_uri(uri)[0]
         assert name
         return name
 
@@ -222,15 +219,15 @@ def test_read_storage_listing_happens_once(cloud_test_catalog, cloud_type):
     dc_dogs = chain.filter(dc.C("file.path").glob("dogs*"))
     dc_cats.union(dc_dogs).save(ds_name)
 
-    lst_ds_name = parse_listing_uri(uri, ctc.session.catalog.client_config)[0]
+    lst_ds_name = parse_listing_uri(uri)[0]
     assert _get_listing_datasets(ctc.session) == [f"{lst_ds_name}@v1.0.0"]
 
 
 def test_read_storage_dependencies(cloud_test_catalog, cloud_type):
     ctc = cloud_test_catalog
     src_uri = ctc.src_uri
     uri = f"{src_uri}/cats"
-    dep_name, _, _ = parse_listing_uri(uri, ctc.catalog.client_config)
+    dep_name, _, _ = parse_listing_uri(uri)
     ds_name = "dep"
     dc.read_storage(uri, session=ctc.session).save(ds_name)
     dependencies = ctc.session.catalog.get_dataset_dependencies(ds_name, "1.0.0")
@@ -244,7 +241,7 @@ def test_persist_not_affects_dependencies(tmp_dir, test_session):
         (tmp_dir / f"file{i}.txt").write_text(f"file{i}")
 
     uri = tmp_dir.as_uri()
-    dep_name, _, _ = parse_listing_uri(uri, test_session.catalog.client_config)
+    dep_name, _, _ = parse_listing_uri(uri)
     chain = dc.read_storage(uri, session=test_session)  # .persist()
     # calling multiple persists to create temp datasets
     chain = chain.persist()

tests/func/test_dataset_query.py

@@ -965,9 +965,7 @@ def test_dataset_dependencies_one_storage_as_dependency(
     ds_name = uuid.uuid4().hex
     catalog = cloud_test_catalog.catalog
     listing = catalog.listings()[0]
-    dep_name, _, _ = parse_listing_uri(
-        cloud_test_catalog.src_uri, catalog.client_config
-    )
+    dep_name, _, _ = parse_listing_uri(cloud_test_catalog.src_uri)
 
     DatasetQuery(cats_dataset.name, catalog=catalog).save(ds_name)
 
@@ -996,9 +994,7 @@ def test_dataset_dependencies_one_registered_dataset_as_dependency(
     catalog = cloud_test_catalog.catalog
     listing = catalog.listings()[0]
 
-    dep_name, _, _ = parse_listing_uri(
-        cloud_test_catalog.src_uri, catalog.client_config
-    )
+    dep_name, _, _ = parse_listing_uri(cloud_test_catalog.src_uri)
 
     DatasetQuery(name=dogs_dataset.name, catalog=catalog).save(ds_name)
 
@@ -1044,9 +1040,7 @@ def test_dataset_dependencies_multiple_direct_dataset_dependencies(
     ds_name = uuid.uuid4().hex
     catalog = cloud_test_catalog.catalog
     listing = catalog.listings()[0]
-    dep_name, _, _ = parse_listing_uri(
-        cloud_test_catalog.src_uri, catalog.client_config
-    )
+    dep_name, _, _ = parse_listing_uri(cloud_test_catalog.src_uri)
 
     dogs = DatasetQuery(name=dogs_dataset.name, version="1.0.0", catalog=catalog)
     cats = DatasetQuery(name=cats_dataset.name, version="1.0.0", catalog=catalog)
@@ -1116,9 +1110,7 @@ def test_dataset_dependencies_multiple_union(
     ds_name = uuid.uuid4().hex
     catalog = cloud_test_catalog.catalog
     listing = catalog.listings()[0]
-    dep_name, _, _ = parse_listing_uri(
-        cloud_test_catalog.src_uri, catalog.client_config
-    )
+    dep_name, _, _ = parse_listing_uri(cloud_test_catalog.src_uri)
 
     dogs = DatasetQuery(name=dogs_dataset.name, version="1.0.0", catalog=catalog)
     cats = DatasetQuery(name=cats_dataset.name, version="1.0.0", catalog=catalog)

tests/func/test_datasets.py

@@ -666,12 +666,12 @@ def test_dataset_storage_dependencies(cloud_test_catalog, cloud_type, indirect):
     session = ctc.session
     catalog = session.catalog
     uri = cloud_test_catalog.src_uri
-    dep_name, _, _ = parse_listing_uri(ctc.src_uri, catalog.client_config)
+    dep_name, _, _ = parse_listing_uri(ctc.src_uri)
 
     ds_name = "some_ds"
     dc.read_storage(uri, session=session).save(ds_name)
 
-    lst_ds_name, _, _ = parse_listing_uri(uri, catalog.client_config)
+    lst_ds_name, _, _ = parse_listing_uri(uri)
     lst_dataset = catalog.metastore.get_dataset(lst_ds_name)
 
     assert [

tests/func/test_listing.py

@@ -36,10 +36,7 @@ def test_listing_generator(cloud_test_catalog, cloud_type):
 )
 def test_parse_listing_uri(cloud_test_catalog, cloud_type):
     ctc = cloud_test_catalog
-    catalog = ctc.catalog
-    dataset_name, listing_uri, listing_path = parse_listing_uri(
-        f"{ctc.src_uri}/dogs", catalog.client_config
-    )
+    dataset_name, listing_uri, listing_path = parse_listing_uri(f"{ctc.src_uri}/dogs")
     assert dataset_name == f"lst__{ctc.src_uri}/dogs/"
     assert listing_uri == f"{ctc.src_uri}/dogs/"
     if cloud_type == "file":
@@ -55,10 +52,7 @@ def test_parse_listing_uri(cloud_test_catalog, cloud_type):
 )
 def test_parse_listing_uri_with_glob(cloud_test_catalog):
     ctc = cloud_test_catalog
-    catalog = ctc.catalog
-    dataset_name, listing_uri, listing_path = parse_listing_uri(
-        f"{ctc.src_uri}/dogs/*", catalog.client_config
-    )
+    dataset_name, listing_uri, listing_path = parse_listing_uri(f"{ctc.src_uri}/dogs/*")
     assert dataset_name == f"lst__{ctc.src_uri}/dogs/"
     assert listing_uri == f"{ctc.src_uri}/dogs"
     assert listing_path == "dogs/*"