datachain 0.14.5__tar.gz → 0.16.0__tar.gz

{datachain-0.14.5 → datachain-0.16.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.11.4'
+    rev: 'v0.11.5'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.14.5/src/datachain.egg-info → datachain-0.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.14.5
+Version: 0.16.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0

{datachain-0.14.5 → datachain-0.16.0}/docs/references/datachain.md

@@ -31,6 +31,10 @@ for examples of how to create a chain.
 
 ::: datachain.lib.dc.values.read_values
 
+::: datachain.lib.dc.database.read_database
+
+::: datachain.lib.dc.database.ConnectionType
+
 ::: datachain.lib.dc.DataChain
 
 ::: datachain.lib.utils.DataChainError

{datachain-0.14.5 → datachain-0.16.0}/examples/multimodal/wds_filtered.py

@@ -27,7 +27,7 @@ filtered = (
         / func.least("laion.json.original_width", "laion.json.original_height")
         < 3.0
     )
-    .save()
+    .persist()
 )
 
 filtered.show(3)

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/__init__.py

@@ -5,8 +5,10 @@ from datachain.lib.dc import (
     DataChain,
     Sys,
     datasets,
+    delete_dataset,
     listings,
     read_csv,
+    read_database,
     read_dataset,
     read_hf,
     read_json,
@@ -61,11 +63,13 @@ __all__ = [
     "VideoFragment",
     "VideoFrame",
     "datasets",
+    "delete_dataset",
     "is_chain_type",
     "listings",
     "metrics",
     "param",
     "read_csv",
+    "read_database",
     "read_dataset",
     "read_hf",
     "read_json",

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/catalog/catalog.py

@@ -776,7 +776,7 @@ class Catalog:
         listing: Optional[bool] = False,
         uuid: Optional[str] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -794,16 +794,16 @@ class Catalog:
             dataset = self.get_dataset(name)
             default_version = dataset.next_version
 
-            if (description or labels) and (
-                dataset.description != description or dataset.labels != labels
+            if (description or attrs) and (
+                dataset.description != description or dataset.attrs != attrs
             ):
                 description = description or dataset.description
-                labels = labels or dataset.labels
+                attrs = attrs or dataset.attrs
 
                 self.update_dataset(
                     dataset,
                     description=description,
-                    labels=labels,
+                    attrs=attrs,
                 )
 
         except DatasetNotFoundError:
@@ -817,7 +817,7 @@ class Catalog:
                 schema=schema,
                 ignore_if_exists=True,
                 description=description,
-                labels=labels,
+                attrs=attrs,
             )
 
         version = version or default_version
@@ -1299,7 +1299,17 @@ class Catalog:
         name: str,
         version: Optional[int] = None,
         force: Optional[bool] = False,
+        studio: Optional[bool] = False,
     ):
+        from datachain.remote.studio import StudioClient
+
+        if studio:
+            client = StudioClient()
+            response = client.rm_dataset(name, version=version, force=force)
+            if not response.ok:
+                raise DataChainError(response.message)
+            return
+
         dataset = self.get_dataset(name)
         if not version and not force:
             raise ValueError(f"Missing dataset version from input for dataset {name}")
@@ -1324,15 +1334,15 @@ class Catalog:
         name: str,
         new_name: Optional[str] = None,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> DatasetRecord:
         update_data = {}
         if new_name:
             update_data["name"] = new_name
         if description is not None:
             update_data["description"] = description
-        if labels is not None:
-            update_data["labels"] = labels  # type: ignore[assignment]
+        if attrs is not None:
+            update_data["attrs"] = attrs  # type: ignore[assignment]
 
         dataset = self.get_dataset(name)
         return self.update_dataset(dataset, **update_data)

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/catalog/loader.py

@@ -1,4 +1,5 @@
 import os
+import sys
 from importlib import import_module
 from typing import TYPE_CHECKING, Any, Optional
 
@@ -15,6 +16,7 @@ METASTORE_ARG_PREFIX = "DATACHAIN_METASTORE_ARG_"
 WAREHOUSE_SERIALIZED = "DATACHAIN__WAREHOUSE"
 WAREHOUSE_IMPORT_PATH = "DATACHAIN_WAREHOUSE"
 WAREHOUSE_ARG_PREFIX = "DATACHAIN_WAREHOUSE_ARG_"
+DISTRIBUTED_IMPORT_PYTHONPATH = "DATACHAIN_DISTRIBUTED_PYTHONPATH"
 DISTRIBUTED_IMPORT_PATH = "DATACHAIN_DISTRIBUTED"
 
 IN_MEMORY_ERROR_MESSAGE = "In-memory is only supported on SQLite"
@@ -100,19 +102,21 @@ def get_warehouse(in_memory: bool = False) -> "AbstractWarehouse":
     return warehouse_class(**warehouse_args)
 
 
-def get_udf_distributor_class() -> type["AbstractUDFDistributor"]:
-    distributed_import_path = os.environ.get(DISTRIBUTED_IMPORT_PATH)
+def get_udf_distributor_class() -> Optional[type["AbstractUDFDistributor"]]:
+    if not (distributed_import_path := os.environ.get(DISTRIBUTED_IMPORT_PATH)):
+        return None
 
-    if not distributed_import_path:
-        raise RuntimeError(
-            f"{DISTRIBUTED_IMPORT_PATH} import path is required "
-            "for distributed UDF processing."
-        )
     # Distributed class paths are specified as (for example): module.classname
     if "." not in distributed_import_path:
         raise RuntimeError(
             f"Invalid {DISTRIBUTED_IMPORT_PATH} import path: {distributed_import_path}"
         )
+
+    # Optional: set the Python path to look for the module
+    distributed_import_pythonpath = os.environ.get(DISTRIBUTED_IMPORT_PYTHONPATH)
+    if distributed_import_pythonpath and distributed_import_pythonpath not in sys.path:
+        sys.path.insert(0, distributed_import_pythonpath)
+
     module_name, _, class_name = distributed_import_path.rpartition(".")
     distributed = import_module(module_name)
     return getattr(distributed, class_name)

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/cli/__init__.py

@@ -149,7 +149,7 @@ def handle_dataset_command(args, catalog):
             args.name,
             new_name=args.new_name,
             description=args.description,
-            labels=args.labels,
+            attrs=args.attrs,
             studio=args.studio,
             local=args.local,
             all=args.all,

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/cli/commands/datasets.py

@@ -154,7 +154,7 @@ def edit_dataset(
     name: str,
     new_name: Optional[str] = None,
     description: Optional[str] = None,
-    labels: Optional[list[str]] = None,
+    attrs: Optional[list[str]] = None,
     studio: bool = False,
     local: bool = False,
     all: bool = True,
@@ -167,9 +167,9 @@ def edit_dataset(
 
     if all or local:
         try:
-            catalog.edit_dataset(name, new_name, description, labels)
+            catalog.edit_dataset(name, new_name, description, attrs)
         except DatasetNotFoundError:
             print("Dataset not found in local", file=sys.stderr)
 
     if (all or studio) and token:
-        edit_studio_dataset(team, name, new_name, description, labels)
+        edit_studio_dataset(team, name, new_name, description, attrs)

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/cli/commands/show.py

@@ -42,8 +42,8 @@ def show(
     print("Name: ", name)
     if dataset.description:
         print("Description: ", dataset.description)
-    if dataset.labels:
-        print("Labels: ", ",".join(dataset.labels))
+    if dataset.attrs:
+        print("Attributes: ", ",".join(dataset.attrs))
     print("\n")
 
     show_records(records, collapse_columns=not no_collapse, hidden_fields=hidden_fields)

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/cli/parser/__init__.py

@@ -217,9 +217,9 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="Dataset description",
     )
     parse_edit_dataset.add_argument(
-        "--labels",
+        "--attrs",
         nargs="+",
-        help="Dataset labels",
+        help="Dataset attributes",
     )
     parse_edit_dataset.add_argument(
         "--studio",

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/data_storage/metastore.py

@@ -120,7 +120,7 @@ class AbstractMetastore(ABC, Serializable):
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
     ) -> DatasetRecord:
         """Creates new dataset."""
 
@@ -326,7 +326,7 @@ class AbstractDBMetastore(AbstractMetastore):
             Column("id", Integer, primary_key=True),
             Column("name", Text, nullable=False),
             Column("description", Text),
-            Column("labels", JSON, nullable=True),
+            Column("attrs", JSON, nullable=True),
             Column("status", Integer, nullable=False),
             Column("feature_schema", JSON, nullable=True),
             Column("created_at", DateTime(timezone=True)),
@@ -521,7 +521,7 @@ class AbstractDBMetastore(AbstractMetastore):
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
         description: Optional[str] = None,
-        labels: Optional[list[str]] = None,
+        attrs: Optional[list[str]] = None,
         **kwargs,  # TODO registered = True / False
     ) -> DatasetRecord:
         """Creates new dataset."""
@@ -538,7 +538,7 @@ class AbstractDBMetastore(AbstractMetastore):
             query_script=query_script,
             schema=json.dumps(schema or {}),
             description=description,
-            labels=json.dumps(labels or []),
+            attrs=json.dumps(attrs or []),
         )
         if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
             # SQLite and PostgreSQL both support 'on_conflict_do_nothing',
@@ -621,7 +621,7 @@ class AbstractDBMetastore(AbstractMetastore):
         dataset_values = {}
         for field, value in kwargs.items():
             if field in self._dataset_fields[1:]:
-                if field in ["labels", "schema"]:
+                if field in ["attrs", "schema"]:
                     values[field] = json.dumps(value) if value else None
                 else:
                     values[field] = value

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/dataset.py

@@ -329,7 +329,7 @@ class DatasetRecord:
     id: int
     name: str
     description: Optional[str]
-    labels: list[str]
+    attrs: list[str]
     schema: dict[str, Union[SQLType, type[SQLType]]]
     feature_schema: dict
     versions: list[DatasetVersion]
@@ -357,7 +357,7 @@ class DatasetRecord:
         id: int,
         name: str,
         description: Optional[str],
-        labels: str,
+        attrs: str,
         status: int,
         feature_schema: Optional[str],
         created_at: datetime,
@@ -387,7 +387,7 @@ class DatasetRecord:
         version_schema: str,
         version_job_id: Optional[str] = None,
     ) -> "DatasetRecord":
-        labels_lst: list[str] = json.loads(labels) if labels else []
+        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 {}
@@ -418,7 +418,7 @@ class DatasetRecord:
             id,
             name,
             description,
-            labels_lst,
+            attrs_lst,
             cls.parse_schema(schema_dct),  # type: ignore[arg-type]
             json.loads(feature_schema) if feature_schema else {},
             [dataset_version],
@@ -562,7 +562,7 @@ class DatasetListRecord:
     id: int
     name: str
     description: Optional[str]
-    labels: list[str]
+    attrs: list[str]
     versions: list[DatasetListVersion]
     created_at: Optional[datetime] = None
 
@@ -572,7 +572,7 @@ class DatasetListRecord:
         id: int,
         name: str,
         description: Optional[str],
-        labels: str,
+        attrs: str,
         created_at: datetime,
         version_id: int,
         version_uuid: str,
@@ -588,7 +588,7 @@ class DatasetListRecord:
         version_query_script: Optional[str],
         version_job_id: Optional[str] = None,
     ) -> "DatasetListRecord":
-        labels_lst: list[str] = json.loads(labels) if labels else []
+        attrs_lst: list[str] = json.loads(attrs) if attrs else []
 
         dataset_version = DatasetListVersion.parse(
             version_id,
@@ -610,7 +610,7 @@ class DatasetListRecord:
             id,
             name,
             description,
-            labels_lst,
+            attrs_lst,
             [dataset_version],
             created_at,
         )

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/lib/convert/values_to_tuples.py

@@ -1,5 +1,6 @@
+import itertools
 from collections.abc import Sequence
-from typing import Any, Union
+from typing import Any, Optional, Union
 
 from datachain.lib.data_model import (
     DataType,
@@ -66,21 +67,29 @@ def values_to_tuples(  # noqa: C901, PLR0912
                     f"signal '{k}' is not present in the output",
                 )
         else:
-            if len_ == 0:
-                raise ValuesToTupleError(ds_name, f"signal '{k}' is empty list")
-
-            first_element = next(iter(v))
-            typ = type(first_element)
-            if not is_chain_type(typ):
-                raise ValuesToTupleError(
-                    ds_name,
-                    f"signal '{k}' has unsupported type '{typ.__name__}'."
-                    f" Please use DataModel types: {DataTypeNames}",
+            # FIXME: Stops as soon as it finds the first non-None value.
+            # If a non-None value appears early, it won't check the remaining items for
+            # `None` values.
+            try:
+                pos, first_not_none_element = next(
+                    itertools.dropwhile(lambda pair: pair[1] is None, enumerate(v))
                 )
-            if isinstance(first_element, list):
-                types_map[k] = list[type(first_element[0])]  # type: ignore[assignment, misc]
+            except StopIteration:
+                typ = str  # default to str if all values are None or has length 0
+                nullable = True
             else:
-                types_map[k] = typ
+                nullable = pos > 0
+                typ = type(first_not_none_element)  # type: ignore[assignment]
+                if not is_chain_type(typ):
+                    raise ValuesToTupleError(
+                        ds_name,
+                        f"signal '{k}' has unsupported type '{typ.__name__}'."
+                        f" Please use DataModel types: {DataTypeNames}",
+                    )
+                if isinstance(first_not_none_element, list):
+                    typ = list[type(first_not_none_element[0])]  # type: ignore[assignment, misc]
+
+            types_map[k] = Optional[typ] if nullable else typ  # type: ignore[assignment]
 
         if length < 0:
             length = len_

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/lib/dataset_info.py

@@ -32,11 +32,28 @@ class DatasetInfo(DataModel):
     metrics: dict[str, Any] = Field(default={})
     error_message: str = Field(default="")
     error_stack: str = Field(default="")
+    attrs: list[str] = Field(default=[])
 
     @property
     def is_temp(self) -> bool:
         return Session.is_temp_dataset(self.name)
 
+    def has_attr(self, attr: str) -> bool:
+        s = attr.split("=")
+        if len(s) == 1:
+            return attr in self.attrs
+
+        name = s[0]
+        value = s[1]
+        for a in self.attrs:
+            s = a.split("=")
+            if value == "*" and s[0] == name:
+                return True
+            if len(s) == 2 and s[0] == name and s[1] == value:
+                return True
+
+        return False
+
     @staticmethod
     def _validate_dict(
         v: Optional[Union[str, dict]],
@@ -83,4 +100,5 @@ class DatasetInfo(DataModel):
             metrics=job.metrics if job else {},
             error_message=version.error_message,
             error_stack=version.error_stack,
+            attrs=dataset.attrs,
         )

{datachain-0.14.5 → datachain-0.16.0}/src/datachain/lib/dc/__init__.py

@@ -1,6 +1,7 @@
 from .csv import read_csv
+from .database import read_database
 from .datachain import C, Column, DataChain
-from .datasets import datasets, read_dataset
+from .datasets import datasets, delete_dataset, read_dataset
 from .hf import read_hf
 from .json import read_json
 from .listings import listings
@@ -19,8 +20,10 @@ __all__ = [
     "DatasetPrepareError",
     "Sys",
     "datasets",
+    "delete_dataset",
     "listings",
     "read_csv",
+    "read_database",
     "read_dataset",
     "read_hf",
     "read_json",

datachain-0.16.0/src/datachain/lib/dc/database.py

@@ -0,0 +1,151 @@
+import contextlib
+import itertools
+import os
+import sqlite3
+from typing import TYPE_CHECKING, Any, Optional, Union
+
+import sqlalchemy
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator, Mapping, Sequence
+
+    import sqlalchemy.orm  # noqa: TC004
+
+    from datachain.lib.data_model import DataType
+    from datachain.query import Session
+
+    from .datachain import DataChain
+
+    ConnectionType = Union[
+        str,
+        sqlalchemy.engine.URL,
+        sqlalchemy.engine.interfaces.Connectable,
+        sqlalchemy.engine.Engine,
+        sqlalchemy.engine.Connection,
+        sqlalchemy.orm.Session,
+        sqlite3.Connection,
+    ]
+
+
+@contextlib.contextmanager
+def _connect(
+    connection: "ConnectionType",
+) -> "Iterator[Union[sqlalchemy.engine.Connection, sqlalchemy.orm.Session]]":
+    import sqlalchemy.orm
+
+    with contextlib.ExitStack() as stack:
+        engine_kwargs = {"echo": bool(os.environ.get("DEBUG_SHOW_SQL_QUERIES"))}
+        if isinstance(connection, (str, sqlalchemy.URL)):
+            engine = sqlalchemy.create_engine(connection, **engine_kwargs)
+            stack.callback(engine.dispose)
+            yield stack.enter_context(engine.connect())
+        elif isinstance(connection, sqlite3.Connection):
+            engine = sqlalchemy.create_engine(
+                "sqlite://", creator=lambda: connection, **engine_kwargs
+            )
+            # do not close the connection, as it is managed by the caller
+            yield engine.connect()
+        elif isinstance(connection, sqlalchemy.Engine):
+            yield stack.enter_context(connection.connect())
+        elif isinstance(connection, (sqlalchemy.Connection, sqlalchemy.orm.Session)):
+            # do not close the connection, as it is managed by the caller
+            yield connection
+        else:
+            raise TypeError(f"Unsupported connection type: {type(connection).__name__}")
+
+
+def _infer_schema(
+    result: "sqlalchemy.engine.Result",
+    to_infer: list[str],
+    infer_schema_length: Optional[int] = 100,
+) -> tuple[list["sqlalchemy.Row"], dict[str, "DataType"]]:
+    from datachain.lib.convert.values_to_tuples import values_to_tuples
+
+    if not to_infer:
+        return [], {}
+
+    rows = list(itertools.islice(result, infer_schema_length))
+    values = {col: [row._mapping[col] for row in rows] for col in to_infer}
+    _, output_schema, _ = values_to_tuples("", **values)
+    return rows, output_schema
+
+
+def read_database(
+    query: Union[str, "sqlalchemy.sql.expression.Executable"],
+    connection: "ConnectionType",
+    params: Union["Sequence[Mapping[str, Any]]", "Mapping[str, Any]", None] = None,
+    *,
+    output: Optional["dict[str, DataType]"] = None,
+    session: Optional["Session"] = None,
+    settings: Optional[dict] = None,
+    in_memory: bool = False,
+    infer_schema_length: Optional[int] = 100,
+) -> "DataChain":
+    """
+    Read the results of a SQL query into a DataChain, using a given database connection.
+
+    Args:
+        query:
+            The SQL query to execute. Can be a raw SQL string or a SQLAlchemy
+            `Executable` object.
+        connection: SQLAlchemy connectable, str, or a sqlite3 connection
+            Using SQLAlchemy makes it possible to use any DB supported by that
+            library. If a DBAPI2 object, only sqlite3 is supported. The user is
+            responsible for engine disposal and connection closure for the
+            SQLAlchemy connectable; str connections are closed automatically.
+        params: Parameters to pass to execute method.
+        output: A dictionary mapping column names to types, used to override the
+            schema inferred from the query results.
+        session: Session to use for the chain.
+        settings: Settings to use for the chain.
+        in_memory: If True, creates an in-memory session. Defaults to False.
+        infer_schema_length:
+            The maximum number of rows to scan for inferring schema.
+            If set to `None`, the full data may be scanned.
+            The rows used for schema inference are stored in memory,
+            so large values can lead to high memory usage.
+            Only applies if the `output` parameter is not set for the given column.
+
+    Examples:
+        Reading from a SQL query against a user-supplied connection:
+        ```python
+        query = "SELECT key, value FROM tbl"
+        chain = dc.read_database(query, connection, output={"value": float})
+        ```
+
+        Load data from a SQLAlchemy driver/engine:
+        ```python
+        from sqlalchemy import create_engine
+        engine = create_engine("postgresql+psycopg://myuser:mypassword@localhost:5432/mydb")
+        chain = dc.read_database("select * from tbl", engine)
+        ```
+
+        Load data from a parameterized SQLAlchemy query:
+        ```python
+        query = "SELECT key, value FROM tbl WHERE value > :value"
+        dc.read_database(query, engine, params={"value": 50})
+        ```
+
+    Notes:
+        This function works with a variety of databases — including, but not limited to,
+        SQLite, DuckDB, PostgreSQL, and Snowflake, provided the appropriate driver is
+        installed.
+    """
+    from datachain.lib.dc.records import read_records
+
+    output = output or {}
+    if isinstance(query, str):
+        query = sqlalchemy.text(query)
+    kw = {"execution_options": {"stream_results": True}}  # use server-side cursors
+    with _connect(connection) as conn, conn.execute(query, params, **kw) as result:
+        cols = result.keys()
+        to_infer = [k for k in cols if k not in output]  # preserve the order
+        rows, inferred_schema = _infer_schema(result, to_infer, infer_schema_length)
+        records = (row._asdict() for row in itertools.chain(rows, result))
+        return read_records(
+            records,
+            session=session,
+            settings=settings,
+            in_memory=in_memory,
+            schema=inferred_schema | output,
+        )