datachain 0.5.0__py3-none-any.whl → 0.6.0__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/data_storage/metastore.py

@@ -50,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")
 
 
@@ -384,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):
     """
@@ -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/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)

datachain/lib/convert/values_to_tuples.py

@@ -4,7 +4,7 @@ from typing import Any, Union
 from datachain.lib.data_model import (
     DataType,
     DataTypeNames,
-    DataValuesType,
+    DataValue,
     is_chain_type,
 )
 from datachain.lib.utils import DataChainParamsError
@@ -20,7 +20,7 @@ class ValuesToTupleError(DataChainParamsError):
 def values_to_tuples(  # noqa: C901, PLR0912
     ds_name: str = "",
     output: Union[None, DataType, Sequence[str], dict[str, DataType]] = None,
-    **fr_map: Sequence[DataValuesType],
+    **fr_map: Sequence[DataValue],
 ) -> tuple[Any, Any, Any]:
     if output:
         if not isinstance(output, (Sequence, str, dict)):

datachain/lib/data_model.py

@@ -18,7 +18,7 @@ StandardType = Union[
 ]
 DataType = Union[type[BaseModel], StandardType]
 DataTypeNames = "BaseModel, int, str, float, bool, list, dict, bytes, datetime"
-DataValuesType = Union[BaseModel, int, str, float, bool, list, dict, bytes, datetime]
+DataValue = Union[BaseModel, int, str, float, bool, list, dict, bytes, datetime]
 
 
 class DataModel(BaseModel):

datachain/lib/dc.py

@@ -16,6 +16,7 @@ from typing import (
     overload,
 )
 
+import orjson
 import pandas as pd
 import sqlalchemy
 from pydantic import BaseModel
@@ -58,9 +59,10 @@ from datachain.query.dataset import (
 from datachain.query.schema import DEFAULT_DELIMITER, Column, DatasetRow
 from datachain.sql.functions import path as pathfunc
 from datachain.telemetry import telemetry
-from datachain.utils import inside_notebook
+from datachain.utils import batched_it, inside_notebook
 
 if TYPE_CHECKING:
+    from pyarrow import DataType as ArrowDataType
     from typing_extensions import Concatenate, ParamSpec, Self
 
     from datachain.lib.hf import HFDatasetType
@@ -71,6 +73,10 @@ C = Column
 
 _T = TypeVar("_T")
 D = TypeVar("D", bound="DataChain")
+UDFObjT = TypeVar("UDFObjT", bound=UDFBase)
+
+
+DEFAULT_PARQUET_CHUNK_SIZE = 100_000
 
 
 def resolve_columns(
@@ -819,7 +825,7 @@ class DataChain:
 
     def gen(
         self,
-        func: Optional[Callable] = None,
+        func: Optional[Union[Callable, Generator]] = None,
         params: Union[None, str, Sequence[str]] = None,
         output: OutputType = None,
         **signal_map,
@@ -931,12 +937,12 @@ class DataChain:
 
     def _udf_to_obj(
         self,
-        target_class: type[UDFBase],
-        func: Optional[Callable],
+        target_class: type[UDFObjT],
+        func: Optional[Union[Callable, UDFObjT]],
         params: Union[None, str, Sequence[str]],
         output: OutputType,
         signal_map,
-    ) -> UDFBase:
+    ) -> UDFObjT:
         is_generator = target_class.is_output_batched
         name = self.name or ""
 
@@ -1019,7 +1025,7 @@ class DataChain:
         The supported functions:
            Numerical:   +, -, *, /, rand(), avg(), count(), func(),
                         greatest(), least(), max(), min(), sum()
-           String:      length(), split()
+           String:      length(), split(), replace(), regexp_replace()
            Filename:    name(), parent(), file_stem(), file_ext()
            Array:       length(), sip_hash_64(), euclidean_distance(),
                         cosine_distance()
@@ -1103,6 +1109,29 @@ class DataChain:
                 rows = (row_factory(db_signals, r) for r in rows)
             yield from rows
 
+    def to_columnar_data_with_names(
+        self, chunk_size: int = DEFAULT_PARQUET_CHUNK_SIZE
+    ) -> tuple[list[str], Iterator[list[list[Any]]]]:
+        """Returns column names and the results as an iterator that provides chunks,
+        with each chunk containing a list of columns, where each column contains a
+        list of the row values for that column in that chunk. Useful for columnar data
+        formats, such as parquet or other OLAP databases.
+        """
+        headers, _ = self._effective_signals_schema.get_headers_with_length()
+        column_names = [".".join(filter(None, header)) for header in headers]
+
+        results_iter = self.collect_flatten()
+
+        def column_chunks() -> Iterator[list[list[Any]]]:
+            for chunk_iter in batched_it(results_iter, chunk_size):
+                columns: list[list[Any]] = [[] for _ in column_names]
+                for row in chunk_iter:
+                    for i, col in enumerate(columns):
+                        col.append(row[i])
+                yield columns
+
+        return column_names, column_chunks()
+
     @overload
     def results(self) -> list[tuple[Any, ...]]: ...
 
@@ -1681,6 +1710,7 @@ class DataChain:
         nrows=None,
         session: Optional[Session] = None,
         settings: Optional[dict] = None,
+        column_types: Optional[dict[str, "Union[str, ArrowDataType]"]] = None,
         **kwargs,
     ) -> "DataChain":
         """Generate chain from csv files.
@@ -1699,6 +1729,9 @@ class DataChain:
             nrows : Optional row limit.
             session : Session to use for the chain.
             settings : Settings to use for the chain.
+            column_types : Dictionary of column names and their corresponding types.
+                It is passed to CSV reader and for each column specified type auto
+                inference is disabled.
 
         Example:
             Reading a csv file:
@@ -1714,6 +1747,15 @@ class DataChain:
         from pandas.io.parsers.readers import STR_NA_VALUES
         from pyarrow.csv import ConvertOptions, ParseOptions, ReadOptions
         from pyarrow.dataset import CsvFileFormat
+        from pyarrow.lib import type_for_alias
+
+        if column_types:
+            column_types = {
+                name: type_for_alias(typ) if isinstance(typ, str) else typ
+                for name, typ in column_types.items()
+            }
+        else:
+            column_types = {}
 
         chain = DataChain.from_storage(
             path, session=session, settings=settings, **kwargs
@@ -1739,7 +1781,9 @@ class DataChain:
         parse_options = ParseOptions(delimiter=delimiter)
         read_options = ReadOptions(column_names=column_names)
         convert_options = ConvertOptions(
-            strings_can_be_null=True, null_values=STR_NA_VALUES
+            strings_can_be_null=True,
+            null_values=STR_NA_VALUES,
+            column_types=column_types,
         )
         format = CsvFileFormat(
             parse_options=parse_options,
@@ -1808,21 +1852,96 @@ class DataChain:
         self,
         path: Union[str, os.PathLike[str], BinaryIO],
         partition_cols: Optional[Sequence[str]] = None,
+        chunk_size: int = DEFAULT_PARQUET_CHUNK_SIZE,
         **kwargs,
     ) -> None:
-        """Save chain to parquet file.
+        """Save chain to parquet file with SignalSchema metadata.
 
         Parameters:
             path : Path or a file-like binary object to save the file.
             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.
         """
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        from datachain.lib.arrow import DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY
+
         _partition_cols = list(partition_cols) if partition_cols else None
-        return self.to_pandas().to_parquet(
-            path,
-            partition_cols=_partition_cols,
-            **kwargs,
+        signal_schema_metadata = orjson.dumps(
+            self._effective_signals_schema.serialize()
         )
 
+        column_names, column_chunks = self.to_columnar_data_with_names(chunk_size)
+
+        parquet_schema = None
+        parquet_writer = None
+        first_chunk = True
+
+        for chunk in column_chunks:
+            # pyarrow infers the best parquet schema from the python types of
+            # the input data.
+            table = pa.Table.from_pydict(
+                dict(zip(column_names, chunk)),
+                schema=parquet_schema,
+            )
+
+            # Preserve any existing metadata, and add the DataChain SignalSchema.
+            existing_metadata = table.schema.metadata or {}
+            merged_metadata = {
+                **existing_metadata,
+                DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY: signal_schema_metadata,
+            }
+            table = table.replace_schema_metadata(merged_metadata)
+            parquet_schema = table.schema
+
+            if _partition_cols:
+                # Write to a partitioned parquet dataset.
+                pq.write_to_dataset(
+                    table,
+                    root_path=path,
+                    partition_cols=_partition_cols,
+                    **kwargs,
+                )
+            else:
+                if first_chunk:
+                    # Write to a single parquet file.
+                    parquet_writer = pq.ParquetWriter(path, parquet_schema, **kwargs)
+                    first_chunk = False
+
+                assert parquet_writer
+                parquet_writer.write_table(table)
+
+        if parquet_writer:
+            parquet_writer.close()
+
+    def to_csv(
+        self,
+        path: Union[str, os.PathLike[str]],
+        delimiter: str = ",",
+        **kwargs,
+    ) -> None:
+        """Save chain to a csv (comma-separated values) file.
+
+        Parameters:
+            path : Path to save the file.
+            delimiter : Delimiter to use for the resulting file.
+        """
+        import csv
+
+        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:
+            writer = csv.writer(f, delimiter=delimiter, **kwargs)
+            writer.writerow(column_names)
+
+            for row in results_iter:
+                writer.writerow(row)
+
     @classmethod
     def from_records(
         cls,

datachain/lib/signal_schema.py

@@ -25,7 +25,7 @@ from typing_extensions import Literal as LiteralEx
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.sql_to_python import sql_to_python
 from datachain.lib.convert.unflatten import unflatten_to_json_pos
-from datachain.lib.data_model import DataModel, DataType
+from datachain.lib.data_model import DataModel, DataType, DataValue
 from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
@@ -110,7 +110,7 @@ class SignalSchema:
     values: dict[str, DataType]
     tree: dict[str, Any]
     setup_func: dict[str, Callable]
-    setup_values: Optional[dict[str, Callable]]
+    setup_values: Optional[dict[str, Any]]
 
     def __init__(
         self,
@@ -333,21 +333,21 @@ class SignalSchema:
                 res[db_name] = python_to_sql(type_)
         return res
 
-    def row_to_objs(self, row: Sequence[Any]) -> list[DataType]:
+    def row_to_objs(self, row: Sequence[Any]) -> list[DataValue]:
         self._init_setup_values()
 
-        objs = []
+        objs: list[DataValue] = []
         pos = 0
         for name, fr_type in self.values.items():
             if self.setup_values and (val := self.setup_values.get(name, None)):
                 objs.append(val)
             elif (fr := ModelStore.to_pydantic(fr_type)) is not None:
                 j, pos = unflatten_to_json_pos(fr, row, pos)
-                objs.append(fr(**j))  # type: ignore[arg-type]
+                objs.append(fr(**j))
             else:
                 objs.append(row[pos])
                 pos += 1
-        return objs  # type: ignore[return-value]
+        return objs
 
     def contains_file(self) -> bool:
         for type_ in self.values.values():