datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/dataset.py

@@ -1,21 +1,14 @@
 import builtins
-import json
 from dataclasses import dataclass, fields
 from datetime import datetime
 from functools import cached_property
-from typing import (
-    Any,
-    NewType,
-    Optional,
-    TypeVar,
-    Union,
-)
+from typing import Any, NewType, TypeVar
 from urllib.parse import urlparse
 
 from packaging.specifiers import SpecifierSet
 from packaging.version import Version
 
-from datachain import semver
+from datachain import json, semver
 from datachain.error import DatasetVersionNotFoundError, InvalidDatasetNameError
 from datachain.namespace import Namespace
 from datachain.project import Project
@@ -43,7 +36,7 @@ DATASET_NAME_REPLACEMENT_CHAR = "_"
 StorageURI = NewType("StorageURI", str)
 
 
-def parse_dataset_uri(uri: str) -> tuple[str, Optional[str]]:
+def parse_dataset_uri(uri: str) -> tuple[str, str | None]:
     """
     Parse dataser uri to extract name and version out of it (if version is defined)
     Example:
@@ -65,7 +58,7 @@ def parse_dataset_uri(uri: str) -> tuple[str, Optional[str]]:
 
 
 def create_dataset_uri(
-    name: str, namespace: str, project: str, version: Optional[str] = None
+    name: str, namespace: str, project: str, version: str | None = None
 ) -> str:
     """
     Creates a dataset uri based on namespace, project, dataset name and optionally
@@ -81,7 +74,7 @@ def create_dataset_uri(
     return uri
 
 
-def parse_dataset_name(name: str) -> tuple[Optional[str], Optional[str], str]:
+def parse_dataset_name(name: str) -> tuple[str | None, str | None, str]:
     """Parses dataset name and returns namespace, project and name"""
     if not name:
         raise InvalidDatasetNameError("Name must be defined to parse it")
@@ -95,6 +88,40 @@ def parse_dataset_name(name: str) -> tuple[Optional[str], Optional[str], str]:
     return namespace_name, project_name, name
 
 
+def parse_schema(ct: dict[str, Any]) -> dict[str, SQLType | type[SQLType]]:
+    """Parse dataset schema from dictionary representation.
+
+    Args:
+        ct: Dictionary with column definitions
+
+    Returns:
+        Dictionary mapping column names to SQL types
+
+    Raises:
+        TypeError: If schema format is invalid
+        ValueError: If column type is not defined or not supported
+    """
+    if not isinstance(ct, dict):
+        raise TypeError("Schema definition must be a dictionary")
+    res = {}
+    for c_name, c_type in ct.items():
+        if not isinstance(c_type, dict):
+            raise TypeError(f"Schema column '{c_name}' type must be a dictionary")
+        if "type" not in c_type:
+            raise ValueError(f"Schema column '{c_name}' type is not defined")
+        if c_type["type"] not in NAME_TYPES_MAPPING:
+            raise ValueError(
+                f"Schema column '{c_name}' type '{c_type['type']}' is not supported"
+            )
+        try:
+            res[c_name] = NAME_TYPES_MAPPING[c_type["type"]].from_dict(c_type)  # type: ignore [attr-defined]
+        except Exception as e:
+            raise ValueError(
+                f"Schema column '{c_name}' type '{c_type['type']}' parsing error: {e}"
+            ) from e
+    return res
+
+
 class DatasetDependencyType:
     DATASET = "dataset"
     STORAGE = "storage"
@@ -111,7 +138,7 @@ class DatasetDependency:
     name: str
     version: str
     created_at: datetime
-    dependencies: list[Optional["DatasetDependency"]]
+    dependencies: list["DatasetDependency | None"]
 
     @property
     def dataset_name(self) -> str:
@@ -131,12 +158,12 @@ class DatasetDependency:
         namespace_name: str,
         project_name: str,
         id: int,
-        dataset_id: Optional[int],
-        dataset_version_id: Optional[int],
-        dataset_name: Optional[str],
-        dataset_version: Optional[str],
-        dataset_version_created_at: Optional[datetime],
-    ) -> Optional["DatasetDependency"]:
+        dataset_id: int | None,
+        dataset_version_id: int | None,
+        dataset_name: str | None,
+        dataset_version: str | None,
+        dataset_version_created_at: datetime | None,
+    ) -> "DatasetDependency | None":
         from datachain.lib.listing import is_listing_dataset
 
         if not dataset_id:
@@ -198,17 +225,17 @@ class DatasetVersion:
     status: int
     feature_schema: dict
     created_at: datetime
-    finished_at: Optional[datetime]
+    finished_at: datetime | None
     error_message: str
     error_stack: str
     script_output: str
-    schema: dict[str, Union[SQLType, type[SQLType]]]
-    num_objects: Optional[int]
-    size: Optional[int]
-    _preview_data: Optional[Union[str, list[dict]]]
+    schema: dict[str, SQLType | type[SQLType]]
+    num_objects: int | None
+    size: int | None
+    _preview_data: str | list[dict] | None
     sources: str = ""
     query_script: str = ""
-    job_id: Optional[str] = None
+    job_id: str | None = None
 
     @classmethod
     def parse(  # noqa: PLR0913
@@ -218,20 +245,25 @@ class DatasetVersion:
         dataset_id: int,
         version: str,
         status: int,
-        feature_schema: Optional[str],
+        feature_schema: str | None,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
         script_output: str,
-        num_objects: Optional[int],
-        size: Optional[int],
-        preview: Optional[Union[str, list[dict]]],
-        schema: dict[str, Union[SQLType, type[SQLType]]],
+        num_objects: int | None,
+        size: int | None,
+        preview: str | list[dict] | None,
+        schema: str | dict[str, SQLType | type[SQLType]],
         sources: str = "",
         query_script: str = "",
-        job_id: Optional[str] = None,
+        job_id: str | None = None,
     ):
+        if isinstance(schema, str):
+            schema_parsed = parse_schema(json.loads(schema) if schema else {})
+        else:
+            schema_parsed = schema
+
         return cls(
             id,
             uuid,
@@ -244,7 +276,7 @@ class DatasetVersion:
             error_message,
             error_stack,
             script_output,
-            schema,
+            schema_parsed,
             num_objects,
             size,
             preview,
@@ -292,7 +324,7 @@ class DatasetVersion:
         }
 
     @cached_property
-    def preview(self) -> Optional[list[dict]]:
+    def preview(self) -> list[dict] | None:
         if isinstance(self._preview_data, str):
             return json.loads(self._preview_data)
         return self._preview_data if self._preview_data else None
@@ -313,13 +345,13 @@ class DatasetListVersion:
     version: str
     status: int
     created_at: datetime
-    finished_at: Optional[datetime]
+    finished_at: datetime | None
     error_message: str
     error_stack: str
-    num_objects: Optional[int]
-    size: Optional[int]
+    num_objects: int | None
+    size: int | None
     query_script: str = ""
-    job_id: Optional[str] = None
+    job_id: str | None = None
 
     @classmethod
     def parse(
@@ -330,13 +362,13 @@ class DatasetListVersion:
         version: str,
         status: int,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
-        num_objects: Optional[int],
-        size: Optional[int],
+        num_objects: int | None,
+        size: int | None,
         query_script: str = "",
-        job_id: Optional[str] = None,
+        job_id: str | None = None,
         **kwargs,
     ):
         return cls(
@@ -368,14 +400,14 @@ class DatasetRecord:
     id: int
     name: str
     project: Project
-    description: Optional[str]
+    description: str | None
     attrs: list[str]
-    schema: dict[str, Union[SQLType, type[SQLType]]]
+    schema: dict[str, SQLType | type[SQLType]]
     feature_schema: dict
     versions: list[DatasetVersion]
     status: int = DatasetStatus.CREATED
-    created_at: Optional[datetime] = None
-    finished_at: Optional[datetime] = None
+    created_at: datetime | None = None
+    finished_at: datetime | None = None
     error_message: str = ""
     error_stack: str = ""
     script_output: str = ""
@@ -385,15 +417,6 @@ class DatasetRecord:
     def __hash__(self):
         return hash(f"{self.id}")
 
-    @staticmethod
-    def parse_schema(
-        ct: dict[str, Any],
-    ) -> dict[str, Union[SQLType, type[SQLType]]]:
-        return {
-            c_name: NAME_TYPES_MAPPING[c_type["type"]].from_dict(c_type)  # type: ignore [attr-defined]
-            for c_name, c_type in ct.items()
-        }
-
     @staticmethod
     def validate_name(name: str) -> None:
         """Throws exception if name has reserved characters"""
@@ -409,23 +432,23 @@ class DatasetRecord:
         namespace_id: int,
         namespace_uuid: str,
         namespace_name: str,
-        namespace_description: Optional[str],
+        namespace_description: str | None,
         namespace_created_at: datetime,
         project_id: int,
         project_uuid: str,
         project_name: str,
-        project_description: Optional[str],
+        project_description: str | None,
         project_created_at: datetime,
         project_namespace_id: int,
         dataset_id: int,
         dataset_project_id: int,
         name: str,
-        description: Optional[str],
+        description: str | None,
         attrs: str,
         status: int,
-        feature_schema: Optional[str],
+        feature_schema: str | None,
         created_at: datetime,
-        finished_at: Optional[datetime],
+        finished_at: datetime | None,
         error_message: str,
         error_stack: str,
         script_output: str,
@@ -437,25 +460,22 @@ class DatasetRecord:
         version_dataset_id: int,
         version: str,
         version_status: int,
-        version_feature_schema: Optional[str],
+        version_feature_schema: str | None,
         version_created_at: datetime,
-        version_finished_at: Optional[datetime],
+        version_finished_at: datetime | None,
         version_error_message: str,
         version_error_stack: str,
         version_script_output: str,
-        version_num_objects: Optional[int],
-        version_size: Optional[int],
-        version_preview: Optional[str],
-        version_sources: Optional[str],
-        version_query_script: Optional[str],
+        version_num_objects: int | None,
+        version_size: int | None,
+        version_preview: str | None,
+        version_sources: str | None,
+        version_query_script: str | None,
         version_schema: str,
-        version_job_id: Optional[str] = None,
+        version_job_id: str | None = None,
     ) -> "DatasetRecord":
         attrs_lst: list[str] = json.loads(attrs) if attrs else []
         schema_dct: dict[str, Any] = json.loads(schema) if schema else {}
-        version_schema_dct: dict[str, str] = (
-            json.loads(version_schema) if version_schema else {}
-        )
 
         namespace = Namespace(
             namespace_id,
@@ -489,7 +509,7 @@ class DatasetRecord:
             version_num_objects,
             version_size,
             version_preview,
-            cls.parse_schema(version_schema_dct),  # type: ignore[arg-type]
+            version_schema,
             version_sources,  # type: ignore[arg-type]
             version_query_script,  # type: ignore[arg-type]
             version_job_id,
@@ -501,7 +521,7 @@ class DatasetRecord:
             project,
             description,
             attrs_lst,
-            cls.parse_schema(schema_dct),  # type: ignore[arg-type]
+            parse_schema(schema_dct),  # type: ignore[arg-type]
             json.loads(feature_schema) if feature_schema else {},
             [dataset_version],
             status,
@@ -527,7 +547,7 @@ class DatasetRecord:
     def full_name(self) -> str:
         return f"{self.project.namespace.name}.{self.project.name}.{self.name}"
 
-    def get_schema(self, version: str) -> dict[str, Union[SQLType, type[SQLType]]]:
+    def get_schema(self, version: str) -> dict[str, SQLType | type[SQLType]]:
         return self.get_version(version).schema if version else self.schema
 
     def update(self, **kwargs):
@@ -619,7 +639,7 @@ class DatasetRecord:
         if not self.versions:
             return "1.0.0"
 
-        major, minor, patch = semver.parse(self.latest_version)
+        major, _, _ = semver.parse(self.latest_version)
         return semver.create(major + 1, 0, 0)
 
     @property
@@ -630,7 +650,7 @@ class DatasetRecord:
         if not self.versions:
             return "1.0.0"
 
-        major, minor, patch = semver.parse(self.latest_version)
+        major, minor, _ = semver.parse(self.latest_version)
         return semver.create(major, minor + 1, 0)
 
     @property
@@ -649,7 +669,7 @@ class DatasetRecord:
         """Returns latest version of a dataset"""
         return max(self.versions).version
 
-    def latest_major_version(self, major: int) -> Optional[str]:
+    def latest_major_version(self, major: int) -> str | None:
         """
         Returns latest specific major version, e.g if dataset has versions:
             - 1.4.1
@@ -664,7 +684,7 @@ class DatasetRecord:
             return None
         return max(versions).version
 
-    def latest_compatible_version(self, version_spec: str) -> Optional[str]:
+    def latest_compatible_version(self, version_spec: str) -> str | None:
         """
         Returns the latest version that matches the given version specifier.
 
@@ -711,10 +731,10 @@ class DatasetListRecord:
     id: int
     name: str
     project: Project
-    description: Optional[str]
+    description: str | None
     attrs: list[str]
     versions: list[DatasetListVersion]
-    created_at: Optional[datetime] = None
+    created_at: datetime | None = None
 
     @classmethod
     def parse(  # noqa: PLR0913
@@ -722,17 +742,17 @@ class DatasetListRecord:
         namespace_id: int,
         namespace_uuid: str,
         namespace_name: str,
-        namespace_description: Optional[str],
+        namespace_description: str | None,
         namespace_created_at: datetime,
         project_id: int,
         project_uuid: str,
         project_name: str,
-        project_description: Optional[str],
+        project_description: str | None,
         project_created_at: datetime,
         project_namespace_id: int,
         dataset_id: int,
         name: str,
-        description: Optional[str],
+        description: str | None,
         attrs: str,
         created_at: datetime,
         version_id: int,
@@ -741,13 +761,13 @@ class DatasetListRecord:
         version: str,
         version_status: int,
         version_created_at: datetime,
-        version_finished_at: Optional[datetime],
+        version_finished_at: datetime | None,
         version_error_message: str,
         version_error_stack: str,
-        version_num_objects: Optional[int],
-        version_size: Optional[int],
-        version_query_script: Optional[str],
-        version_job_id: Optional[str] = None,
+        version_num_objects: int | None,
+        version_size: int | None,
+        version_query_script: str | None,
+        version_job_id: str | None = None,
     ) -> "DatasetListRecord":
         attrs_lst: list[str] = json.loads(attrs) if attrs else []
 
@@ -824,7 +844,7 @@ class DatasetListRecord:
         from datachain.client import Client
 
         # TODO refactor and maybe remove method in
-        # https://github.com/iterative/datachain/issues/318
+        # https://github.com/datachain-ai/datachain/issues/318
         return Client.is_data_source_uri(self.name) or self.name.startswith(
             LISTING_PREFIX
         )