datachain 0.3.12__py3-none-any.whl → 0.3.14__py3-none-any.whl

datachain/lib/file.py

@@ -118,7 +118,6 @@ class File(DataModel):
     is_latest: bool = Field(default=True)
     last_modified: datetime = Field(default=TIME_ZERO)
     location: Optional[Union[dict, list[dict]]] = Field(default=None)
-    vtype: str = Field(default="")
 
     _datachain_column_types: ClassVar[dict[str, Any]] = {
         "source": String,
@@ -129,7 +128,6 @@ class File(DataModel):
         "is_latest": Boolean,
         "last_modified": DateTime,
         "location": JSON,
-        "vtype": String,
     }
 
     _unique_id_keys: ClassVar[list[str]] = [
@@ -139,7 +137,6 @@ class File(DataModel):
         "etag",
         "version",
         "is_latest",
-        "vtype",
         "location",
         "last_modified",
     ]

datachain/lib/listing.py

@@ -30,8 +30,7 @@ def list_bucket(uri: str, client_config=None) -> Callable:
         config = client_config or {}
         client, path = Client.parse_url(uri, None, **config)  # type: ignore[arg-type]
         for entries in iter_over_async(client.scandir(path.rstrip("/")), get_loop()):
-            for entry in entries:
-                yield entry.to_file(client.uri)
+            yield from entries
 
     return list_func
 

datachain/lib/model_store.py

@@ -1,6 +1,6 @@
 import inspect
 import logging
-from typing import ClassVar, Optional
+from typing import Any, ClassVar, Optional
 
 from pydantic import BaseModel
 
@@ -69,7 +69,7 @@ class ModelStore:
             del cls.store[fr.__name__][version]
 
     @staticmethod
-    def is_pydantic(val):
+    def is_pydantic(val: Any) -> bool:
         return (
             not hasattr(val, "__origin__")
             and inspect.isclass(val)

datachain/lib/pytorch.py

@@ -7,6 +7,7 @@ from torch import float32
 from torch.distributed import get_rank, get_world_size
 from torch.utils.data import IterableDataset, get_worker_info
 from torchvision.transforms import v2
+from tqdm import tqdm
 
 from datachain.catalog import Catalog, get_catalog
 from datachain.lib.dc import DataChain
@@ -93,33 +94,38 @@ class PytorchDataset(IterableDataset):
         if self.num_samples > 0:
             ds = ds.sample(self.num_samples)
         ds = ds.chunk(total_rank, total_workers)
-        for row_features in ds.collect():
-            row = []
-            for fr in row_features:
-                if hasattr(fr, "read"):
-                    row.append(fr.read())  # type: ignore[unreachable]
-                else:
-                    row.append(fr)
-            # Apply transforms
-            if self.transform:
-                try:
-                    if isinstance(self.transform, v2.Transform):
-                        row = self.transform(row)
+        desc = f"Parsed PyTorch dataset for rank={total_rank} worker"
+        with tqdm(desc=desc, unit=" rows") as pbar:
+            for row_features in ds.collect():
+                row = []
+                for fr in row_features:
+                    if hasattr(fr, "read"):
+                        row.append(fr.read())  # type: ignore[unreachable]
+                    else:
+                        row.append(fr)
+                # Apply transforms
+                if self.transform:
+                    try:
+                        if isinstance(self.transform, v2.Transform):
+                            row = self.transform(row)
+                        for i, val in enumerate(row):
+                            if isinstance(val, Image.Image):
+                                row[i] = self.transform(val)
+                    except ValueError:
+                        logger.warning(
+                            "Skipping transform due to unsupported data types."
+                        )
+                        self.transform = None
+                if self.tokenizer:
                     for i, val in enumerate(row):
-                        if isinstance(val, Image.Image):
-                            row[i] = self.transform(val)
-                except ValueError:
-                    logger.warning("Skipping transform due to unsupported data types.")
-                    self.transform = None
-            if self.tokenizer:
-                for i, val in enumerate(row):
-                    if isinstance(val, str) or (
-                        isinstance(val, list) and isinstance(val[0], str)
-                    ):
-                        row[i] = convert_text(
-                            val, self.tokenizer, self.tokenizer_kwargs
-                        ).squeeze(0)  # type: ignore[union-attr]
-            yield row
+                        if isinstance(val, str) or (
+                            isinstance(val, list) and isinstance(val[0], str)
+                        ):
+                            row[i] = convert_text(
+                                val, self.tokenizer, self.tokenizer_kwargs
+                            ).squeeze(0)  # type: ignore[union-attr]
+                yield row
+                pbar.update(1)
 
     @staticmethod
     def get_rank_and_workers() -> tuple[int, int]:

datachain/lib/signal_schema.py

@@ -4,11 +4,14 @@ from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from datetime import datetime
 from inspect import isclass
-from typing import (
+from typing import (  # noqa: UP035
     TYPE_CHECKING,
     Annotated,
     Any,
     Callable,
+    Dict,
+    Final,
+    List,
     Literal,
     Optional,
     Union,
@@ -42,8 +45,13 @@ NAMES_TO_TYPES = {
     "dict": dict,
     "bytes": bytes,
     "datetime": datetime,
-    "Literal": Literal,
+    "Final": Final,
     "Union": Union,
+    "Optional": Optional,
+    "List": list,
+    "Dict": dict,
+    "Literal": Any,
+    "Any": Any,
 }
 
 
@@ -146,35 +154,11 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def _get_name_original_type(fr_type: type) -> tuple[str, type]:
-        """Returns the name of and the original type for the given type,
-        based on whether the type is Optional or not."""
-        orig = get_origin(fr_type)
-        args = get_args(fr_type)
-        # Check if fr_type is Optional
-        if orig == Union and len(args) == 2 and (type(None) in args):
-            fr_type = args[0]
-            orig = get_origin(fr_type)
-        if orig in (Literal, LiteralEx):
-            # Literal has no __name__ in Python 3.9
-            type_name = "Literal"
-        elif orig == Union:
-            # Union also has no __name__ in Python 3.9
-            type_name = "Union"
-        else:
-            type_name = str(fr_type.__name__)  # type: ignore[union-attr]
-        return type_name, fr_type
-
-    @staticmethod
-    def serialize_custom_model_fields(
-        name: str, fr: type, custom_types: dict[str, Any]
+    def _serialize_custom_model_fields(
+        version_name: str, fr: type[BaseModel], custom_types: dict[str, Any]
     ) -> str:
         """This serializes any custom type information to the provided custom_types
-        dict, and returns the name of the type provided."""
-        if hasattr(fr, "__origin__") or not issubclass(fr, BaseModel):
-            # Don't store non-feature types.
-            return name
-        version_name = ModelStore.get_name(fr)
+        dict, and returns the name of the type serialized."""
         if version_name in custom_types:
             # This type is already stored in custom_types.
             return version_name
@@ -183,37 +167,102 @@ class SignalSchema:
             field_type = info.annotation
             # All fields should be typed.
             assert field_type
-            field_type_name, field_type = SignalSchema._get_name_original_type(
-                field_type
-            )
-            # Serialize this type to custom_types if it is a custom type as well.
-            fields[field_name] = SignalSchema.serialize_custom_model_fields(
-                field_type_name, field_type, custom_types
-            )
+            fields[field_name] = SignalSchema._serialize_type(field_type, custom_types)
         custom_types[version_name] = fields
         return version_name
 
+    @staticmethod
+    def _serialize_type(fr: type, custom_types: dict[str, Any]) -> str:
+        """Serialize a given type to a string, including automatic ModelStore
+        registration, and save this type and subtypes to custom_types as well."""
+        subtypes: list[Any] = []
+        type_name = SignalSchema._type_to_str(fr, subtypes)
+        # Iterate over all subtypes (includes the input type).
+        for st in subtypes:
+            if st is None or not ModelStore.is_pydantic(st):
+                continue
+            # Register and save feature types.
+            ModelStore.register(st)
+            st_version_name = ModelStore.get_name(st)
+            if st is fr:
+                # If the main type is Pydantic, then use the ModelStore version name.
+                type_name = st_version_name
+            # Save this type to custom_types.
+            SignalSchema._serialize_custom_model_fields(
+                st_version_name, st, custom_types
+            )
+        return type_name
+
     def serialize(self) -> dict[str, Any]:
         signals: dict[str, Any] = {}
         custom_types: dict[str, Any] = {}
         for name, fr_type in self.values.items():
-            if (fr := ModelStore.to_pydantic(fr_type)) is not None:
-                ModelStore.register(fr)
-                signals[name] = ModelStore.get_name(fr)
-                type_name, fr_type = SignalSchema._get_name_original_type(fr)
-            else:
-                type_name, fr_type = SignalSchema._get_name_original_type(fr_type)
-                signals[name] = type_name
-            self.serialize_custom_model_fields(type_name, fr_type, custom_types)
+            signals[name] = self._serialize_type(fr_type, custom_types)
         if custom_types:
             signals["_custom_types"] = custom_types
         return signals
 
     @staticmethod
-    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:
+    def _split_subtypes(type_name: str) -> list[str]:
+        """This splits a list of subtypes, including proper square bracket handling."""
+        start = 0
+        depth = 0
+        subtypes = []
+        for i, c in enumerate(type_name):
+            if c == "[":
+                depth += 1
+            elif c == "]":
+                if depth == 0:
+                    raise TypeError(
+                        "Extra closing square bracket when parsing subtype list"
+                    )
+                depth -= 1
+            elif c == "," and depth == 0:
+                subtypes.append(type_name[start:i].strip())
+                start = i + 1
+        if depth > 0:
+            raise TypeError("Unclosed square bracket when parsing subtype list")
+        subtypes.append(type_name[start:].strip())
+        return subtypes
+
+    @staticmethod
+    def _resolve_type(type_name: str, custom_types: dict[str, Any]) -> Optional[type]:  # noqa: PLR0911
         """Convert a string-based type back into a python type."""
+        type_name = type_name.strip()
+        if not type_name:
+            raise TypeError("Type cannot be empty")
+        if type_name == "NoneType":
+            return None
+
+        bracket_idx = type_name.find("[")
+        subtypes: Optional[tuple[Optional[type], ...]] = None
+        if bracket_idx > -1:
+            if bracket_idx == 0:
+                raise TypeError("Type cannot start with '['")
+            close_bracket_idx = type_name.rfind("]")
+            if close_bracket_idx == -1:
+                raise TypeError("Unclosed square bracket when parsing type")
+            if close_bracket_idx < bracket_idx:
+                raise TypeError("Square brackets are out of order when parsing type")
+            if close_bracket_idx == bracket_idx + 1:
+                raise TypeError("Empty square brackets when parsing type")
+            subtype_names = SignalSchema._split_subtypes(
+                type_name[bracket_idx + 1 : close_bracket_idx]
+            )
+            # Types like Union require the parameters to be a tuple of types.
+            subtypes = tuple(
+                SignalSchema._resolve_type(st, custom_types) for st in subtype_names
+            )
+            type_name = type_name[:bracket_idx].strip()
+
         fr = NAMES_TO_TYPES.get(type_name)
         if fr:
+            if subtypes:
+                if len(subtypes) == 1:
+                    # Types like Optional require there to be only one argument.
+                    return fr[subtypes[0]]  # type: ignore[index]
+                # Other types like Union require the parameters to be a tuple of types.
+                return fr[subtypes]  # type: ignore[index]
             return fr  # type: ignore[return-value]
 
         model_name, version = ModelStore.parse_name_version(type_name)
@@ -228,7 +277,14 @@ class SignalSchema:
                 for field_name, field_type_str in fields.items()
             }
             return create_feature_model(type_name, fields)
-        return None
+        # This can occur if a third-party or custom type is used, which is not available
+        # when deserializing.
+        warnings.warn(
+            f"Could not resolve type: '{type_name}'.",
+            SignalSchemaWarning,
+            stacklevel=2,
+        )
+        return Any  # type: ignore[return-value]
 
     @staticmethod
     def deserialize(schema: dict[str, Any]) -> "SignalSchema":
@@ -242,9 +298,14 @@ class SignalSchema:
                 # This entry is used as a lookup for custom types,
                 # and is not an actual field.
                 continue
+            if not isinstance(type_name, str):
+                raise SignalSchemaError(
+                    f"cannot deserialize '{type_name}': "
+                    "serialized types must be a string"
+                )
             try:
                 fr = SignalSchema._resolve_type(type_name, custom_types)
-                if fr is None:
+                if fr is Any:
                     # Skip if the type is not found, so all data can be displayed.
                     warnings.warn(
                         f"In signal '{signal}': "
@@ -258,7 +319,7 @@ class SignalSchema:
                 raise SignalSchemaError(
                     f"cannot deserialize '{signal}': {err}"
                 ) from err
-            signals[signal] = fr
+            signals[signal] = fr  # type: ignore[assignment]
 
         return SignalSchema(signals)
 
@@ -509,31 +570,58 @@ class SignalSchema:
         return self.values.pop(name)
 
     @staticmethod
-    def _type_to_str(type_):  # noqa: PLR0911
+    def _type_to_str(type_: Optional[type], subtypes: Optional[list] = None) -> str:  # noqa: PLR0911
+        """Convert a type to a string-based representation."""
+        if type_ is None:
+            return "NoneType"
+
         origin = get_origin(type_)
 
         if origin == Union:
             args = get_args(type_)
-            formatted_types = ", ".join(SignalSchema._type_to_str(arg) for arg in args)
+            formatted_types = ", ".join(
+                SignalSchema._type_to_str(arg, subtypes) for arg in args
+            )
             return f"Union[{formatted_types}]"
         if origin == Optional:
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0])
+            type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"Optional[{type_str}]"
-        if origin is list:
+        if origin in (list, List):  # noqa: UP006
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0])
+            type_str = SignalSchema._type_to_str(args[0], subtypes)
             return f"list[{type_str}]"
-        if origin is dict:
+        if origin in (dict, Dict):  # noqa: UP006
             args = get_args(type_)
-            type_str = SignalSchema._type_to_str(args[0]) if len(args) > 0 else ""
-            vals = f", {SignalSchema._type_to_str(args[1])}" if len(args) > 1 else ""
+            type_str = (
+                SignalSchema._type_to_str(args[0], subtypes) if len(args) > 0 else ""
+            )
+            vals = (
+                f", {SignalSchema._type_to_str(args[1], subtypes)}"
+                if len(args) > 1
+                else ""
+            )
             return f"dict[{type_str}{vals}]"
         if origin == Annotated:
             args = get_args(type_)
-            return SignalSchema._type_to_str(args[0])
-        if origin in (Literal, LiteralEx):
+            return SignalSchema._type_to_str(args[0], subtypes)
+        if origin in (Literal, LiteralEx) or type_ in (Literal, LiteralEx):
             return "Literal"
+        if Any in (origin, type_):
+            return "Any"
+        if Final in (origin, type_):
+            return "Final"
+        if subtypes is not None:
+            # Include this type in the list of all subtypes, if requested.
+            subtypes.append(type_)
+        if not hasattr(type_, "__name__"):
+            # This can happen for some third-party or custom types, mostly on Python 3.9
+            warnings.warn(
+                f"Unable to determine name of type '{type_}'.",
+                SignalSchemaWarning,
+                stacklevel=2,
+            )
+            return "Any"
         return type_.__name__
 
     @staticmethod

datachain/listing.py

@@ -9,7 +9,8 @@ from sqlalchemy import Column
 from sqlalchemy.sql import func
 from tqdm import tqdm
 
-from datachain.node import DirType, Entry, Node, NodeWithPath
+from datachain.lib.file import File
+from datachain.node import DirType, Node, NodeWithPath
 from datachain.sql.functions import path as pathfunc
 from datachain.utils import suffix_to_number
 
@@ -80,16 +81,13 @@ class Listing:
             finally:
                 fetch_listing.insert_entries_done()
 
-    def insert_entry(self, entry: Entry) -> None:
-        self.warehouse.insert_rows(
-            self.dataset_rows.get_table(),
-            self.warehouse.prepare_entries(self.client.uri, [entry]),
-        )
+    def insert_entry(self, entry: File) -> None:
+        self.insert_entries([entry])
 
-    def insert_entries(self, entries: Iterable[Entry]) -> None:
+    def insert_entries(self, entries: Iterable[File]) -> None:
         self.warehouse.insert_rows(
             self.dataset_rows.get_table(),
-            self.warehouse.prepare_entries(self.client.uri, entries),
+            self.warehouse.prepare_entries(entries),
         )
 
     def insert_entries_done(self) -> None:
@@ -104,7 +102,7 @@ class Listing:
         return self.warehouse.get_node_by_path(self.dataset_rows, path)
 
     def ls_path(self, node, fields):
-        if node.vtype == "tar" or node.dir_type == DirType.TAR_ARCHIVE:
+        if node.location or node.dir_type == DirType.TAR_ARCHIVE:
             return self.warehouse.select_node_fields_by_parent_path_tar(
                 self.dataset_rows, node.path, fields
             )
@@ -235,7 +233,7 @@ class Listing:
         return self.warehouse.size(self.dataset_rows, node, count_files)
 
     def subtree_files(self, node: Node, sort=None):
-        if node.dir_type == DirType.TAR_ARCHIVE or node.vtype != "":
+        if node.dir_type == DirType.TAR_ARCHIVE or node.location:
             include_subobjects = True
         else:
             include_subobjects = False

datachain/node.py

@@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, Optional
 import attrs
 
 from datachain.cache import UniqueId
-from datachain.lib.file import File
 from datachain.storage import StorageURI
 from datachain.utils import TIME_ZERO, time_to_str
 
@@ -49,18 +48,15 @@ class DirTypeGroup:
 class Node:
     sys__id: int = 0
     sys__rand: int = 0
-    vtype: str = ""
-    dir_type: Optional[int] = None
     path: str = ""
     etag: str = ""
     version: Optional[str] = None
     is_latest: bool = True
     last_modified: Optional[datetime] = None
     size: int = 0
-    owner_name: str = ""
-    owner_id: str = ""
     location: Optional[str] = None
     source: StorageURI = StorageURI("")
+    dir_type: int = DirType.FILE
 
     @property
     def is_dir(self) -> bool:
@@ -113,7 +109,6 @@ class Node:
             version=self.version or "",
             etag=self.etag,
             is_latest=self.is_latest,
-            vtype=self.vtype,
             location=self.location,
             last_modified=self.last_modified or TIME_ZERO,
         )
@@ -143,66 +138,6 @@ class Node:
         return split[0]
 
 
-@attrs.define
-class Entry:
-    vtype: str = ""
-    dir_type: Optional[int] = None
-    path: str = ""
-    etag: str = ""
-    version: str = ""
-    is_latest: bool = True
-    last_modified: Optional[datetime] = None
-    size: int = 0
-    owner_name: str = ""
-    owner_id: str = ""
-    location: Optional[str] = None
-
-    @property
-    def is_dir(self) -> bool:
-        return self.dir_type == DirType.DIR
-
-    @classmethod
-    def from_dir(cls, path: str, **kwargs) -> "Entry":
-        return cls(dir_type=DirType.DIR, path=path, **kwargs)
-
-    @classmethod
-    def from_file(cls, path: str, **kwargs) -> "Entry":
-        return cls(dir_type=DirType.FILE, path=path, **kwargs)
-
-    @classmethod
-    def root(cls):
-        return cls(dir_type=DirType.DIR)
-
-    @property
-    def full_path(self) -> str:
-        if self.is_dir and self.path:
-            return self.path + "/"
-        return self.path
-
-    @property
-    def name(self):
-        return self.path.rsplit("/", 1)[-1]
-
-    @property
-    def parent(self):
-        split = self.path.rsplit("/", 1)
-        if len(split) <= 1:
-            return ""
-        return split[0]
-
-    def to_file(self, source: str) -> File:
-        return File(
-            source=source,
-            path=self.path,
-            size=self.size,
-            version=self.version,
-            etag=self.etag,
-            is_latest=self.is_latest,
-            last_modified=self.last_modified,
-            location=self.location,
-        )
-
-
 def get_path(parent: str, name: str):
     return f"{parent}/{name}" if parent else name
 
@@ -229,9 +164,9 @@ class NodeWithPath:
 TIME_FMT = "%Y-%m-%d %H:%M"
 
 
-def long_line_str(name: str, timestamp: Optional[datetime], owner: str) -> str:
+def long_line_str(name: str, timestamp: Optional[datetime]) -> str:
     if timestamp is None:
         time = "-"
     else:
         time = timestamp.strftime(TIME_FMT)
-    return f"{owner: <19} {time: <19} {name}"
+    return f"{time: <19} {name}"

datachain/query/builtins.py

@@ -22,10 +22,6 @@ def load_tar(raw):
         C.source,
         C.path,
         C.size,
-        C.vtype,
-        C.dir_type,
-        C.owner_name,
-        C.owner_id,
         C.is_latest,
         C.last_modified,
         C.version,
@@ -38,10 +34,6 @@ def index_tar(
     source,
     parent_path,
     size,
-    vtype,
-    dir_type,
-    owner_name,
-    owner_id,
     is_latest,
     last_modified,
     version,
@@ -53,10 +45,6 @@ def index_tar(
         source=source,
         path=parent_path,
         size=size,
-        vtype=vtype,
-        dir_type=dir_type,
-        owner_name=owner_name,
-        owner_id=owner_id,
         is_latest=bool(is_latest),
         last_modified=last_modified,
         version=version,
@@ -70,7 +58,6 @@ def index_tar(
                 source=source,
                 path=full_path,
                 size=info.size,
-                vtype="tar",
                 location={
                     "vtype": "tar",
                     "offset": info.offset_data,
@@ -81,7 +68,6 @@ def index_tar(
                         "version": version,
                         "size": size,
                         "etag": etag,
-                        "vtype": "",
                         "location": None,
                     },
                 },

datachain/query/schema.py

@@ -9,7 +9,7 @@ import attrs
 import sqlalchemy as sa
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 
-from datachain.sql.types import JSON, Boolean, DateTime, Int, Int64, SQLType, String
+from datachain.sql.types import JSON, Boolean, DateTime, Int64, SQLType, String
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
@@ -222,10 +222,6 @@ class DatasetRow:
         "path": String,
         "size": Int64,
         "location": JSON,
-        "vtype": String,
-        "dir_type": Int,
-        "owner_name": String,
-        "owner_id": String,
         "is_latest": Boolean,
         "last_modified": DateTime,
         "version": String,
@@ -238,10 +234,6 @@ class DatasetRow:
         source: str = "",
         size: int = 0,
         location: Optional[dict[str, Any]] = None,
-        vtype: str = "",
-        dir_type: int = 0,
-        owner_name: str = "",
-        owner_id: str = "",
         is_latest: bool = True,
         last_modified: Optional[datetime] = None,
         version: str = "",
@@ -251,10 +243,7 @@ class DatasetRow:
         str,
         int,
         Optional[str],
-        str,
         int,
-        str,
-        str,
         bool,
         datetime,
         str,
@@ -271,10 +260,6 @@ class DatasetRow:
             path,
             size,
             location,
-            vtype,
-            dir_type,
-            owner_name,
-            owner_id,
             is_latest,
             last_modified,
             version,