datachain 0.6.1__py3-none-any.whl → 0.6.2__py3-none-any.whl

datachain/lib/dc.py

@@ -1,6 +1,8 @@
 import copy
 import os
+import os.path
 import re
+import sys
 from collections.abc import Iterator, Sequence
 from functools import wraps
 from typing import (
@@ -23,6 +25,9 @@ from pydantic import BaseModel
 from sqlalchemy.sql.functions import GenericFunction
 from sqlalchemy.sql.sqltypes import NullType
 
+from datachain.client import Client
+from datachain.client.local import FileClient
+from datachain.dataset import DatasetRecord
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.values_to_tuples import values_to_tuples
 from datachain.lib.data_model import DataModel, DataType, dict_to_data_model
@@ -31,9 +36,6 @@ from datachain.lib.file import ArrowRow, File, get_file_type
 from datachain.lib.file import ExportPlacement as FileExportPlacement
 from datachain.lib.func import Func
 from datachain.lib.listing import (
-    is_listing_dataset,
-    is_listing_expired,
-    is_listing_subset,
     list_bucket,
     ls,
     parse_listing_uri,
@@ -51,7 +53,7 @@ from datachain.query.dataset import DatasetQuery, PartitionByType
 from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
 from datachain.sql.functions import path as pathfunc
 from datachain.telemetry import telemetry
-from datachain.utils import batched_it, inside_notebook
+from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
 
 if TYPE_CHECKING:
     from pyarrow import DataType as ArrowDataType
@@ -287,6 +289,13 @@ class DataChain:
         """Version of the underlying dataset, if there is one."""
         return self._query.version
 
+    @property
+    def dataset(self) -> Optional[DatasetRecord]:
+        """Underlying dataset, if there is one."""
+        if not self.name:
+            return None
+        return self.session.catalog.get_dataset(self.name)
+
     def __or__(self, other: "Self") -> "Self":
         """Return `self.union(other)`."""
         return self.union(other)
@@ -367,6 +376,47 @@ class DataChain:
         self.signals_schema |= signals_schema
         return self
 
+    @classmethod
+    def parse_uri(
+        cls, uri: str, session: Session, update: bool = False
+    ) -> tuple[str, str, str, bool]:
+        """Returns correct listing dataset name that must be used for saving listing
+        operation. It takes into account existing listings and reusability of those.
+        It also returns boolean saying if returned dataset name is reused / already
+        exists or not, and it returns correct listing path that should be used to find
+        rows based on uri.
+        """
+        catalog = session.catalog
+        cache = catalog.cache
+        client_config = catalog.client_config
+
+        client = Client.get_client(uri, cache, **client_config)
+        ds_name, list_uri, list_path = parse_listing_uri(uri, cache, client_config)
+        listing = None
+
+        listings = [
+            ls
+            for ls in catalog.listings()
+            if not ls.is_expired and ls.contains(ds_name)
+        ]
+
+        if listings:
+            if update:
+                # choosing the smallest possible one to minimize update time
+                listing = sorted(listings, key=lambda ls: len(ls.name))[0]
+            else:
+                # no need to update, choosing the most recent one
+                listing = sorted(listings, key=lambda ls: ls.created_at)[-1]
+
+        if isinstance(client, FileClient) and listing and listing.name != ds_name:
+            # For local file system we need to fix listing path / prefix
+            # if we are reusing existing listing
+            list_path = f'{ds_name.strip("/").removeprefix(listing.name)}/{list_path}'
+
+        ds_name = listing.name if listing else ds_name
+
+        return ds_name, list_uri, list_path, bool(listing)
+
     @classmethod
     def from_storage(
         cls,
@@ -401,25 +451,15 @@ class DataChain:
         file_type = get_file_type(type)
 
         client_config = {"anon": True} if anon else None
-
         session = Session.get(session, client_config=client_config, in_memory=in_memory)
+        cache = session.catalog.cache
+        client_config = session.catalog.client_config
 
-        list_dataset_name, list_uri, list_path = parse_listing_uri(
-            uri, session.catalog.cache, session.catalog.client_config
+        list_ds_name, list_uri, list_path, list_ds_exists = cls.parse_uri(
+            uri, session, update=update
         )
-        need_listing = True
-
-        for ds in cls.listings(session=session, in_memory=in_memory).collect("listing"):
-            if (
-                not is_listing_expired(ds.created_at)  # type: ignore[union-attr]
-                and is_listing_subset(ds.name, list_dataset_name)  # type: ignore[union-attr]
-                and not update
-            ):
-                need_listing = False
-                list_dataset_name = ds.name  # type: ignore[union-attr]
 
-        if need_listing:
-            # caching new listing to special listing dataset
+        if update or not list_ds_exists:
             (
                 cls.from_records(
                     DataChain.DEFAULT_FILE_RECORD,
@@ -428,17 +468,13 @@ class DataChain:
                     in_memory=in_memory,
                 )
                 .gen(
-                    list_bucket(
-                        list_uri,
-                        session.catalog.cache,
-                        client_config=session.catalog.client_config,
-                    ),
+                    list_bucket(list_uri, cache, client_config=client_config),
                     output={f"{object_name}": File},
                 )
-                .save(list_dataset_name, listing=True)
+                .save(list_ds_name, listing=True)
             )
 
-        dc = cls.from_dataset(list_dataset_name, session=session, settings=settings)
+        dc = cls.from_dataset(list_ds_name, session=session, settings=settings)
         dc.signals_schema = dc.signals_schema.mutate({f"{object_name}": file_type})
 
         return ls(dc, list_path, recursive=recursive, object_name=object_name)
@@ -665,19 +701,11 @@ class DataChain:
         session = Session.get(session, in_memory=in_memory)
         catalog = kwargs.get("catalog") or session.catalog
 
-        listings = [
-            ListingInfo.from_models(d, v, j)
-            for d, v, j in catalog.list_datasets_versions(
-                include_listing=True, **kwargs
-            )
-            if is_listing_dataset(d.name)
-        ]
-
         return cls.from_values(
             session=session,
             in_memory=in_memory,
             output={object_name: ListingInfo},
-            **{object_name: listings},  # type: ignore[arg-type]
+            **{object_name: catalog.listings()},  # type: ignore[arg-type]
         )
 
     def print_json_schema(  # type: ignore[override]
@@ -1006,6 +1034,9 @@ class DataChain:
         """Group rows by specified set of signals and return new signals
         with aggregated values.
 
+        The supported functions:
+           count(), sum(), avg(), min(), max(), any_value(), collect(), concat()
+
         Example:
             ```py
             chain = chain.group_by(
@@ -1069,13 +1100,22 @@ class DataChain:
            Filename:    name(), parent(), file_stem(), file_ext()
            Array:       length(), sip_hash_64(), euclidean_distance(),
                         cosine_distance()
+           Window:      row_number(), rank(), dense_rank(), first()
 
         Example:
         ```py
          dc.mutate(
-                area=Column("image.height") * Column("image.width"),
-                extension=file_ext(Column("file.name")),
-                dist=cosine_distance(embedding_text, embedding_image)
+            area=Column("image.height") * Column("image.width"),
+            extension=file_ext(Column("file.name")),
+            dist=cosine_distance(embedding_text, embedding_image)
+        )
+        ```
+
+        Window function example:
+        ```py
+        window = func.window(partition_by="file.parent", order_by="file.size")
+        dc.mutate(
+            row_number=func.row_number().over(window),
         )
         ```
 
@@ -1086,20 +1126,12 @@ class DataChain:
         Example:
         ```py
          dc.mutate(
-                newkey=Column("oldkey")
+            newkey=Column("oldkey")
         )
         ```
         """
-        existing_columns = set(self.signals_schema.values.keys())
-        for col_name in kwargs:
-            if col_name in existing_columns:
-                raise DataChainColumnError(
-                    col_name,
-                    "Cannot modify existing column with mutate(). "
-                    "Use a different name for the new column.",
-                )
         for col_name, expr in kwargs.items():
-            if not isinstance(expr, Column) and isinstance(expr.type, NullType):
+            if not isinstance(expr, (Column, Func)) and isinstance(expr.type, NullType):
                 raise DataChainColumnError(
                     col_name, f"Cannot infer type with expression {expr}"
                 )
@@ -1111,6 +1143,9 @@ class DataChain:
                 # renaming existing column
                 for signal in schema.db_signals(name=value.name, as_columns=True):
                     mutated[signal.name.replace(value.name, name, 1)] = signal  # type: ignore[union-attr]
+            elif isinstance(value, Func):
+                # adding new signal
+                mutated[name] = value.get_column(schema)
             else:
                 # adding new signal
                 mutated[name] = value
@@ -1887,21 +1922,48 @@ class DataChain:
         path: Union[str, os.PathLike[str], BinaryIO],
         partition_cols: Optional[Sequence[str]] = None,
         chunk_size: int = DEFAULT_PARQUET_CHUNK_SIZE,
+        fs_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ) -> None:
         """Save chain to parquet file with SignalSchema metadata.
 
         Parameters:
-            path : Path or a file-like binary object to save the file.
+            path : Path or a file-like binary object to save the file. This supports
+                local paths as well as remote paths, such as s3:// or hf:// with fsspec.
             partition_cols : Column names by which to partition the dataset.
             chunk_size : The chunk size of results to read and convert to columnar
                 data, to avoid running out of memory.
+            fs_kwargs : Optional kwargs to pass to the fsspec filesystem, used only for
+                write, for fsspec-type URLs, such as s3:// or hf:// when
+                provided as the destination path.
         """
         import pyarrow as pa
         import pyarrow.parquet as pq
 
         from datachain.lib.arrow import DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY
 
+        fsspec_fs = None
+
+        if isinstance(path, str) and "://" in path:
+            from datachain.client.fsspec import Client
+
+            fs_kwargs = {
+                **self._query.catalog.client_config,
+                **(fs_kwargs or {}),
+            }
+
+            client = Client.get_implementation(path)
+
+            if path.startswith("file://"):
+                # pyarrow does not handle file:// uris, and needs a direct path instead.
+                from urllib.parse import urlparse
+
+                path = urlparse(path).path
+                if sys.platform == "win32":
+                    path = os.path.normpath(path.lstrip("/"))
+
+            fsspec_fs = client.create_fs(**fs_kwargs)
+
         _partition_cols = list(partition_cols) if partition_cols else None
         signal_schema_metadata = orjson.dumps(
             self._effective_signals_schema.serialize()
@@ -1936,12 +1998,15 @@ class DataChain:
                     table,
                     root_path=path,
                     partition_cols=_partition_cols,
+                    filesystem=fsspec_fs,
                     **kwargs,
                 )
             else:
                 if first_chunk:
                     # Write to a single parquet file.
-                    parquet_writer = pq.ParquetWriter(path, parquet_schema, **kwargs)
+                    parquet_writer = pq.ParquetWriter(
+                        path, parquet_schema, filesystem=fsspec_fs, **kwargs
+                    )
                     first_chunk = False
 
                 assert parquet_writer
@@ -1954,28 +2019,122 @@ class DataChain:
         self,
         path: Union[str, os.PathLike[str]],
         delimiter: str = ",",
+        fs_kwargs: Optional[dict[str, Any]] = None,
         **kwargs,
     ) -> None:
         """Save chain to a csv (comma-separated values) file.
 
         Parameters:
-            path : Path to save the file.
+            path : Path to save the file. This supports local paths as well as
+                remote paths, such as s3:// or hf:// with fsspec.
             delimiter : Delimiter to use for the resulting file.
+            fs_kwargs : Optional kwargs to pass to the fsspec filesystem, used only for
+                write, for fsspec-type URLs, such as s3:// or hf:// when
+                provided as the destination path.
         """
         import csv
 
+        opener = open
+
+        if isinstance(path, str) and "://" in path:
+            from datachain.client.fsspec import Client
+
+            fs_kwargs = {
+                **self._query.catalog.client_config,
+                **(fs_kwargs or {}),
+            }
+
+            client = Client.get_implementation(path)
+
+            fsspec_fs = client.create_fs(**fs_kwargs)
+
+            opener = fsspec_fs.open
+
         headers, _ = self._effective_signals_schema.get_headers_with_length()
         column_names = [".".join(filter(None, header)) for header in headers]
 
         results_iter = self.collect_flatten()
 
-        with open(path, "w", newline="") as f:
+        with opener(path, "w", newline="") as f:
             writer = csv.writer(f, delimiter=delimiter, **kwargs)
             writer.writerow(column_names)
 
             for row in results_iter:
                 writer.writerow(row)
 
+    def to_json(
+        self,
+        path: Union[str, os.PathLike[str]],
+        fs_kwargs: Optional[dict[str, Any]] = None,
+        include_outer_list: bool = True,
+    ) -> None:
+        """Save chain to a JSON file.
+
+        Parameters:
+            path : Path to save the file. This supports local paths as well as
+                remote paths, such as s3:// or hf:// with fsspec.
+            fs_kwargs : Optional kwargs to pass to the fsspec filesystem, used only for
+                write, for fsspec-type URLs, such as s3:// or hf:// when
+                provided as the destination path.
+            include_outer_list : Sets whether to include an outer list for all rows.
+                Setting this to True makes the file valid JSON, while False instead
+                writes in the JSON lines format.
+        """
+        opener = open
+
+        if isinstance(path, str) and "://" in path:
+            from datachain.client.fsspec import Client
+
+            fs_kwargs = {
+                **self._query.catalog.client_config,
+                **(fs_kwargs or {}),
+            }
+
+            client = Client.get_implementation(path)
+
+            fsspec_fs = client.create_fs(**fs_kwargs)
+
+            opener = fsspec_fs.open
+
+        headers, _ = self._effective_signals_schema.get_headers_with_length()
+        headers = [list(filter(None, header)) for header in headers]
+
+        is_first = True
+
+        with opener(path, "wb") as f:
+            if include_outer_list:
+                # This makes the file JSON instead of JSON lines.
+                f.write(b"[\n")
+            for row in self.collect_flatten():
+                if not is_first:
+                    if include_outer_list:
+                        # This makes the file JSON instead of JSON lines.
+                        f.write(b",\n")
+                    else:
+                        f.write(b"\n")
+                else:
+                    is_first = False
+                f.write(orjson.dumps(row_to_nested_dict(headers, row)))
+            if include_outer_list:
+                # This makes the file JSON instead of JSON lines.
+                f.write(b"\n]\n")
+
+    def to_jsonl(
+        self,
+        path: Union[str, os.PathLike[str]],
+        fs_kwargs: Optional[dict[str, Any]] = None,
+    ) -> None:
+        """Save chain to a JSON lines file.
+
+        Parameters:
+            path : Path to save the file. This supports local paths as well as
+                remote paths, such as s3:// or hf:// with fsspec.
+            fs_kwargs : Optional kwargs to pass to the fsspec filesystem, used only for
+                write, for fsspec-type URLs, such as s3:// or hf:// when
+                provided as the destination path.
+        """
+        self.to_json(path, fs_kwargs, include_outer_list=False)
+
     @classmethod
     def from_records(
         cls,

datachain/lib/func/__init__.py

@@ -1,5 +1,18 @@
-from .aggregate import any_value, avg, collect, concat, count, max, min, sum
-from .func import Func
+from .aggregate import (
+    any_value,
+    avg,
+    collect,
+    concat,
+    count,
+    dense_rank,
+    first,
+    max,
+    min,
+    rank,
+    row_number,
+    sum,
+)
+from .func import Func, window
 
 __all__ = [
     "Func",
@@ -8,7 +21,12 @@ __all__ = [
     "collect",
     "concat",
     "count",
+    "dense_rank",
+    "first",
     "max",
     "min",
+    "rank",
+    "row_number",
     "sum",
+    "window",
 ]