datachain 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

datachain/catalog/catalog.py

@@ -988,6 +988,14 @@ class Catalog:
         schema = {
             c.name: c.type.to_dict() for c in columns if isinstance(c.type, SQLType)
         }
+
+        job_id = job_id or os.getenv("DATACHAIN_JOB_ID")
+        if not job_id:
+            from datachain.query.session import Session
+
+            session = Session.get(catalog=self)
+            job_id = session.job_id
+
         dataset = self.metastore.create_dataset_version(
             dataset,
             version,

datachain/cli.py

@@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union
 
 import shtab
 
-from datachain import utils
+from datachain import Session, utils
 from datachain.cli_utils import BooleanOptionalAction, CommaSeparatedArgs, KeyValueArgs
 from datachain.lib.dc import DataChain
 from datachain.telemetry import telemetry
@@ -770,7 +770,8 @@ def show(
     show_records(records, collapse_columns=not no_collapse)
     if schema and dataset_version.feature_schema:
         print("\nSchema:")
-        dc = DataChain(name=name, version=version, catalog=catalog)
+        session = Session.get(catalog=catalog)
+        dc = DataChain.from_dataset(name=name, version=version, session=session)
         dc.print_schema()
 
 

datachain/data_storage/metastore.py

@@ -15,7 +15,6 @@ from uuid import uuid4
 from sqlalchemy import (
     JSON,
     BigInteger,
-    Boolean,
     Column,
     DateTime,
     ForeignKey,
@@ -51,7 +50,6 @@ if TYPE_CHECKING:
     from datachain.data_storage import AbstractIDGenerator, schema
     from datachain.data_storage.db_engine import DatabaseEngine
 
-
 logger = logging.getLogger("datachain")
 
 
@@ -228,7 +226,7 @@ class AbstractMetastore(ABC, Serializable):
         self,
         dataset: DatasetRecord,
         version: int,
-        status: int = DatasetStatus.CREATED,
+        status: int,
         sources: str = "",
         feature_schema: Optional[dict] = None,
         query_script: str = "",
@@ -385,6 +383,11 @@ class AbstractMetastore(ABC, Serializable):
     ) -> None:
         """Set the status of the given job and dataset."""
 
+    @abstractmethod
+    def get_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
+        """Returns dataset names and versions for the job."""
+        raise NotImplementedError
+
 
 class AbstractDBMetastore(AbstractMetastore):
     """
@@ -448,7 +451,6 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("name", Text, nullable=False),
             Column("description", Text),
             Column("labels", JSON, nullable=True),
-            Column("shadow", Boolean, nullable=False),
             Column("status", Integer, nullable=False),
             Column("feature_schema", JSON, nullable=True),
             Column("created_at", DateTime(timezone=True)),
@@ -481,8 +483,11 @@ class AbstractDBMetastore(AbstractMetastore):
                 nullable=False,
             ),
             Column("version", Integer, nullable=False),
-            # adding default for now until we fully remove shadow datasets
-            Column("status", Integer, nullable=False, default=DatasetStatus.COMPLETE),
+            Column(
+                "status",
+                Integer,
+                nullable=False,
+            ),
             Column("feature_schema", JSON, nullable=True),
             Column("created_at", DateTime(timezone=True)),
             Column("finished_at", DateTime(timezone=True)),
@@ -969,7 +974,6 @@ class AbstractDBMetastore(AbstractMetastore):
         # TODO abstract this method and add registered = True based on kwargs
         query = self._datasets_insert().values(
             name=name,
-            shadow=False,
             status=status,
             feature_schema=json.dumps(feature_schema or {}),
             created_at=datetime.now(timezone.utc),
@@ -992,7 +996,7 @@ class AbstractDBMetastore(AbstractMetastore):
         self,
         dataset: DatasetRecord,
         version: int,
-        status: int = DatasetStatus.CREATED,
+        status: int,
         sources: str = "",
         feature_schema: Optional[dict] = None,
         query_script: str = "",
@@ -1018,7 +1022,7 @@ class AbstractDBMetastore(AbstractMetastore):
         query = self._datasets_versions_insert().values(
             dataset_id=dataset.id,
             version=version,
-            status=status,  # for now until we remove shadow datasets
+            status=status,
             feature_schema=json.dumps(feature_schema or {}),
             created_at=created_at or datetime.now(timezone.utc),
             finished_at=finished_at,
@@ -1519,3 +1523,18 @@ class AbstractDBMetastore(AbstractMetastore):
                 .values(status=dataset_status)
             )
             self.db.execute(query, conn=conn)  # type: ignore[attr-defined]
+
+    def get_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
+        """Returns dataset names and versions for the job."""
+        dv = self._datasets_versions
+        ds = self._datasets
+
+        join_condition = dv.c.dataset_id == ds.c.id
+
+        query = (
+            self._datasets_versions_select(ds.c.name, dv.c.version)
+            .select_from(dv.join(ds, join_condition))
+            .where(dv.c.job_id == job_id)
+        )
+
+        return list(self.db.execute(query))

datachain/data_storage/sqlite.py

@@ -15,6 +15,7 @@ from typing import (
 )
 
 import sqlalchemy
+from packaging import version
 from sqlalchemy import MetaData, Table, UniqueConstraint, exists, select
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
@@ -153,7 +154,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
                 import sys
 
-                db.set_trace_callback(sys.stderr.write)
+                db.set_trace_callback(lambda stmt: print(stmt, file=sys.stderr))
 
             load_usearch_extension(db)
 
@@ -345,45 +346,36 @@ class SQLiteIDGenerator(AbstractDBIDGenerator):
     def get_next_ids(self, uri: str, count: int) -> range:
         """Returns a range of IDs for the given URI."""
 
-        # NOTE: we can't use RETURNING clause here because it is only available
-        # in sqlalchemy v2, see
-        # https://github.com/sqlalchemy/sqlalchemy/issues/6195#issuecomment-1248700677
-        # After we upgrade to sqlalchemy v2, we can use the following code,
-        # leaving fallback to the current implementation for older versions of SQLite,
-        # which is still supported, for example, in Ubuntu 20.04 LTS (Focal Fossa),
-        # where SQLite version 3.31.1 is used.
-
-        # sqlite_version = version.parse(sqlite3.sqlite_version)
-        # if sqlite_version >= version.parse("3.35.0"):
-        #     # RETURNING is supported on SQLite 3.35.0 (2021-03-12) or newer
-        #     stmt = (
-        #         sqlite.insert(self._table)
-        #         .values(uri=uri, last_id=count)
-        #         .on_conflict_do_update(
-        #             index_elements=["uri"],
-        #             set_={"last_id": self._table.c.last_id + count},
-        #         )
-        #         .returning(self._table.c.last_id)
-        #     )
-        #     last_id = self._db.execute(stmt).fetchone()[0]
-        # else:
-        #     (fallback to the current implementation with a transaction)
-
-        # Transactions ensure no concurrency conflicts
-        with self._db.transaction() as conn:
-            # UPSERT syntax was added to SQLite with version 3.24.0 (2018-06-04).
-            stmt_ins = (
+        sqlite_version = version.parse(sqlite3.sqlite_version)
+        is_returning_supported = sqlite_version >= version.parse("3.35.0")
+        if is_returning_supported:
+            stmt = (
                 sqlite.insert(self._table)
                 .values(uri=uri, last_id=count)
                 .on_conflict_do_update(
                     index_elements=["uri"],
                     set_={"last_id": self._table.c.last_id + count},
                 )
+                .returning(self._table.c.last_id)
             )
-            self._db.execute(stmt_ins, conn=conn)
+            last_id = self._db.execute(stmt).fetchone()[0]
+        else:
+            # Older versions of SQLite are still the default under Ubuntu LTS,
+            # e.g. Ubuntu 20.04 LTS (Focal Fossa) uses 3.31.1
+            # Transactions ensure no concurrency conflicts
+            with self._db.transaction() as conn:
+                stmt_ins = (
+                    sqlite.insert(self._table)
+                    .values(uri=uri, last_id=count)
+                    .on_conflict_do_update(
+                        index_elements=["uri"],
+                        set_={"last_id": self._table.c.last_id + count},
+                    )
+                )
+                self._db.execute(stmt_ins, conn=conn)
 
-            stmt_sel = select(self._table.c.last_id).where(self._table.c.uri == uri)
-            last_id = self._db.execute(stmt_sel, conn=conn).fetchone()[0]
+                stmt_sel = select(self._table.c.last_id).where(self._table.c.uri == uri)
+                last_id = self._db.execute(stmt_sel, conn=conn).fetchone()[0]
 
         return range(last_id - count + 1, last_id + 1)
 

datachain/data_storage/warehouse.py

@@ -919,9 +919,7 @@ class AbstractWarehouse(ABC, Serializable):
     def is_temp_table_name(self, name: str) -> bool:
         """Returns if the given table name refers to a temporary
         or no longer needed table."""
-        return name.startswith(
-            (self.TMP_TABLE_NAME_PREFIX, self.UDF_TABLE_NAME_PREFIX, "ds_shadow_")
-        ) or name.endswith("_shadow")
+        return name.startswith((self.TMP_TABLE_NAME_PREFIX, self.UDF_TABLE_NAME_PREFIX))
 
     def get_temp_table_names(self) -> list[str]:
         return [

datachain/dataset.py

@@ -267,7 +267,6 @@ class DatasetRecord:
     name: str
     description: Optional[str]
     labels: list[str]
-    shadow: bool
     schema: dict[str, Union[SQLType, type[SQLType]]]
     feature_schema: dict
     versions: list[DatasetVersion]
@@ -296,7 +295,6 @@ class DatasetRecord:
         name: str,
         description: Optional[str],
         labels: str,
-        shadow: int,
         status: int,
         feature_schema: Optional[str],
         created_at: datetime,
@@ -356,7 +354,6 @@ class DatasetRecord:
             name,
             description,
             labels_lst,
-            bool(shadow),
             cls.parse_schema(schema_dct),  # type: ignore[arg-type]
             json.loads(feature_schema) if feature_schema else {},
             [dataset_version],

datachain/lib/arrow.py

@@ -1,8 +1,9 @@
 import re
 from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
+import orjson
 import pyarrow as pa
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm import tqdm
@@ -10,6 +11,7 @@ from tqdm import tqdm
 from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import ArrowRow, File
 from datachain.lib.model_store import ModelStore
+from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
@@ -20,6 +22,9 @@ if TYPE_CHECKING:
     from datachain.lib.dc import DataChain
 
 
+DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY = b"DataChain SignalSchema"
+
+
 class ArrowGenerator(Generator):
     def __init__(
         self,
@@ -61,28 +66,35 @@ class ArrowGenerator(Generator):
                 path, filesystem=file.get_fs(), schema=self.input_schema, **self.kwargs
             )
         hf_schema = _get_hf_schema(ds.schema)
+        use_datachain_schema = (
+            bool(ds.schema.metadata)
+            and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in ds.schema.metadata
+        )
         index = 0
         with tqdm(desc="Parsed by pyarrow", unit=" rows") as pbar:
             for record_batch in ds.to_batches():
                 for record in record_batch.to_pylist():
-                    vals = list(record.values())
-                    if self.output_schema:
-                        fields = self.output_schema.model_fields
-                        vals_dict = {}
-                        for i, ((field, field_info), val) in enumerate(
-                            zip(fields.items(), vals)
-                        ):
-                            anno = field_info.annotation
-                            if hf_schema:
-                                from datachain.lib.hf import convert_feature
-
-                                feat = list(hf_schema[0].values())[i]
-                                vals_dict[field] = convert_feature(val, feat, anno)
-                            elif ModelStore.is_pydantic(anno):
-                                vals_dict[field] = anno(**val)  # type: ignore[misc]
-                            else:
-                                vals_dict[field] = val
-                        vals = [self.output_schema(**vals_dict)]
+                    if use_datachain_schema and self.output_schema:
+                        vals = [_nested_model_instantiate(record, self.output_schema)]
+                    else:
+                        vals = list(record.values())
+                        if self.output_schema:
+                            fields = self.output_schema.model_fields
+                            vals_dict = {}
+                            for i, ((field, field_info), val) in enumerate(
+                                zip(fields.items(), vals)
+                            ):
+                                anno = field_info.annotation
+                                if hf_schema:
+                                    from datachain.lib.hf import convert_feature
+
+                                    feat = list(hf_schema[0].values())[i]
+                                    vals_dict[field] = convert_feature(val, feat, anno)
+                                elif ModelStore.is_pydantic(anno):
+                                    vals_dict[field] = anno(**val)  # type: ignore[misc]
+                                else:
+                                    vals_dict[field] = val
+                            vals = [self.output_schema(**vals_dict)]
                     if self.source:
                         kwargs: dict = self.kwargs
                         # Can't serialize CsvFileFormat; may lose formatting options.
@@ -113,6 +125,9 @@ def schema_to_output(schema: pa.Schema, col_names: Optional[Sequence[str]] = Non
         )
     if not col_names:
         col_names = schema.names
+    signal_schema = _get_datachain_schema(schema)
+    if signal_schema:
+        return signal_schema.values
     columns = _convert_col_names(col_names)  # type: ignore[arg-type]
     hf_schema = _get_hf_schema(schema)
     if hf_schema:
@@ -197,3 +212,33 @@ def _get_hf_schema(
         features = schema_from_arrow(schema)
         return features, get_output_schema(features)
     return None
+
+
+def _get_datachain_schema(schema: "pa.Schema") -> Optional[SignalSchema]:
+    """Return a restored SignalSchema from parquet metadata, if any is found."""
+    if schema.metadata and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in schema.metadata:
+        serialized_signal_schema = orjson.loads(
+            schema.metadata[DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY]
+        )
+        return SignalSchema.deserialize(serialized_signal_schema)
+    return None
+
+
+def _nested_model_instantiate(
+    column_values: dict[str, Any], model: type["BaseModel"], prefix: str = ""
+) -> "BaseModel":
+    """Instantiate the given model and all sub-models/fields based on the provided
+    column values."""
+    vals_dict = {}
+    for field, field_info in model.model_fields.items():
+        anno = field_info.annotation
+        cur_path = f"{prefix}.{field}" if prefix else field
+        if ModelStore.is_pydantic(anno):
+            vals_dict[field] = _nested_model_instantiate(
+                column_values,
+                anno,  # type: ignore[arg-type]
+                prefix=cur_path,
+            )
+        elif cur_path in column_values:
+            vals_dict[field] = column_values[cur_path]
+    return model(**vals_dict)