datachain 0.14.4__py3-none-any.whl → 0.15.0__py3-none-any.whl

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/catalog/catalog.py

@@ -580,15 +580,13 @@ class Catalog:
         source: str,
         update=False,
         client_config=None,
-        object_name="file",
+        column="file",
         skip_indexing=False,
     ) -> tuple[Optional["Listing"], "Client", str]:
         from datachain import read_storage
         from datachain.listing import Listing
 
-        read_storage(
-            source, session=self.session, update=update, object_name=object_name
-        ).exec()
+        read_storage(source, session=self.session, update=update, column=column).exec()
 
         list_ds_name, list_uri, list_path, _ = get_listing(
             source, self.session, update=update
@@ -602,7 +600,7 @@ class Catalog:
                 self.warehouse.clone(),
                 client,
                 dataset_name=list_ds_name,
-                object_name=object_name,
+                column=column,
             )
 
         return lst, client, list_path
@@ -1301,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}")

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/data_storage/schema.py

@@ -30,8 +30,8 @@ if TYPE_CHECKING:
 DEFAULT_DELIMITER = "__"
 
 
-def col_name(name: str, object_name: str = "file") -> str:
-    return f"{object_name}{DEFAULT_DELIMITER}{name}"
+def col_name(name: str, column: str = "file") -> str:
+    return f"{column}{DEFAULT_DELIMITER}{name}"
 
 
 def dedup_columns(columns: Iterable[sa.Column]) -> list[sa.Column]:
@@ -84,19 +84,19 @@ def convert_rows_custom_column_types(
 
 
 class DirExpansion:
-    def __init__(self, object_name: str):
-        self.object_name = object_name
+    def __init__(self, column: str):
+        self.column = column
 
-    def col_name(self, name: str, object_name: Optional[str] = None) -> str:
-        object_name = object_name or self.object_name
-        return col_name(name, object_name)
+    def col_name(self, name: str, column: Optional[str] = None) -> str:
+        column = column or self.column
+        return col_name(name, column)
 
-    def c(self, query, name: str, object_name: Optional[str] = None) -> str:
-        return getattr(query.c, self.col_name(name, object_name=object_name))
+    def c(self, query, name: str, column: Optional[str] = None) -> str:
+        return getattr(query.c, self.col_name(name, column=column))
 
     def base_select(self, q):
         return sa.select(
-            self.c(q, "id", object_name="sys"),
+            self.c(q, "id", column="sys"),
             false().label(self.col_name("is_dir")),
             self.c(q, "source"),
             self.c(q, "path"),
@@ -153,12 +153,12 @@ class DataTable:
         name: str,
         engine: "DatabaseEngine",
         column_types: Optional[dict[str, SQLType]] = None,
-        object_name: str = "file",
+        column: str = "file",
     ):
         self.name: str = name
         self.engine = engine
         self.column_types: dict[str, SQLType] = column_types or {}
-        self.object_name = object_name
+        self.column = column
 
     @staticmethod
     def copy_column(
@@ -224,18 +224,16 @@ class DataTable:
     def columns(self) -> "ReadOnlyColumnCollection[str, sa.Column[Any]]":
         return self.table.columns
 
-    def col_name(self, name: str, object_name: Optional[str] = None) -> str:
-        object_name = object_name or self.object_name
-        return col_name(name, object_name)
+    def col_name(self, name: str, column: Optional[str] = None) -> str:
+        column = column or self.column
+        return col_name(name, column)
 
-    def without_object(
-        self, column_name: str, object_name: Optional[str] = None
-    ) -> str:
-        object_name = object_name or self.object_name
-        return column_name.removeprefix(f"{object_name}{DEFAULT_DELIMITER}")
+    def without_object(self, column_name: str, column: Optional[str] = None) -> str:
+        column = column or self.column
+        return column_name.removeprefix(f"{column}{DEFAULT_DELIMITER}")
 
-    def c(self, name: str, object_name: Optional[str] = None):
-        return getattr(self.columns, self.col_name(name, object_name=object_name))
+    def c(self, name: str, column: Optional[str] = None):
+        return getattr(self.columns, self.col_name(name, column=column))
 
     @property
     def table(self) -> "sa.Table":
@@ -275,7 +273,7 @@ class DataTable:
         ]
 
     def dir_expansion(self):
-        return DirExpansion(self.object_name)
+        return DirExpansion(self.column)
 
 
 PARTITION_COLUMN_ID = "partition_id"

datachain/data_storage/sqlite.py

@@ -489,7 +489,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self, dataset: DatasetRecord, version: int
     ) -> list[StorageURI]:
         dr = self.dataset_rows(dataset, version)
-        query = dr.select(dr.c("source", object_name="file")).distinct()
+        query = dr.select(dr.c("source", column="file")).distinct()
         cur = self.db.cursor()
         cur.row_factory = sqlite3.Row  # type: ignore[assignment]
 

datachain/data_storage/warehouse.py

@@ -179,7 +179,7 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         dataset: DatasetRecord,
         version: Optional[int] = None,
-        object_name: str = "file",
+        column: str = "file",
     ):
         version = version or dataset.latest_version
 
@@ -188,7 +188,7 @@ class AbstractWarehouse(ABC, Serializable):
             table_name,
             self.db,
             dataset.get_schema(version),
-            object_name=object_name,
+            column=column,
         )
 
     @property
@@ -487,7 +487,7 @@ class AbstractWarehouse(ABC, Serializable):
         dataset_rows: "DataTable",
         path_list: list[str],
         glob_name: str,
-        object_name="file",
+        column="file",
     ) -> Iterator[Node]:
         """Finds all Nodes that correspond to GLOB like path pattern."""
         dr = dataset_rows
@@ -521,7 +521,7 @@ class AbstractWarehouse(ABC, Serializable):
         de = dr.dir_expansion()
         q = de.query(
             dr.select().where(dr.c("is_latest") == true()).subquery(),
-            object_name=dr.object_name,
+            column=dr.column,
         ).subquery()
         q = self.expand_query(de, q, dr)
 
@@ -597,12 +597,10 @@ class AbstractWarehouse(ABC, Serializable):
             with_default(dr.c("is_latest")),
             dr.c("last_modified"),
             with_default(dr.c("size")),
-            with_default(dr.c("rand", object_name="sys")),
+            with_default(dr.c("rand", column="sys")),
             dr.c("location"),
             de.c(q, "source"),
-        ).select_from(
-            q.outerjoin(dr.table, q.c.sys__id == dr.c("id", object_name="sys"))
-        )
+        ).select_from(q.outerjoin(dr.table, q.c.sys__id == dr.c("id", column="sys")))
 
     def get_node_by_path(self, dataset_rows: "DataTable", path: str) -> Node:
         """Gets node that corresponds to some path"""

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/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/lib/dc/csv.py

@@ -21,7 +21,7 @@ def read_csv(
     delimiter: Optional[str] = None,
     header: bool = True,
     output: OutputType = None,
-    object_name: str = "",
+    column: str = "",
     model_name: str = "",
     source: bool = True,
     nrows=None,
@@ -42,7 +42,7 @@ def read_csv(
         output : Dictionary or feature class defining column names and their
             corresponding types. List of column names is also accepted, in which
             case types will be inferred.
-        object_name : Created object column name.
+        column : Created column name.
         model_name : Generated model name.
         source : Whether to include info about the source file.
         nrows : Optional row limit.
@@ -119,7 +119,7 @@ def read_csv(
     )
     return chain.parse_tabular(
         output=output,
-        object_name=object_name,
+        column=column,
         model_name=model_name,
         source=source,
         nrows=nrows,

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,
+        )

datachain/lib/dc/datachain.py

@@ -133,7 +133,7 @@ class DataChain:
                 .choices[0]
                 .message.content,
             )
-            .save()
+            .persist()
         )
 
         try:
@@ -357,7 +357,7 @@ class DataChain:
         self,
         col: str,
         model_name: Optional[str] = None,
-        object_name: Optional[str] = None,
+        column: Optional[str] = None,
         schema_sample_size: int = 1,
     ) -> "DataChain":
         """Explodes a column containing JSON objects (dict or str DataChain type) into
@@ -368,7 +368,7 @@ class DataChain:
             col: the name of the column containing JSON to be exploded.
             model_name: optional generated model name.  By default generates the name
                 automatically.
-            object_name: optional generated object column name. By default generates the
+            column: optional generated column name. By default generates the
                 name automatically.
             schema_sample_size: the number of rows to use for inferring the schema of
                 the JSON (in case some fields are optional and it's not enough to
@@ -406,10 +406,10 @@ class DataChain:
             )
             return model.model_validate(json_dict)
 
-        if not object_name:
-            object_name = f"{col}_expl"
+        if not column:
+            column = f"{col}_expl"
 
-        return self.map(json_to_model, params=col, output={object_name: model})
+        return self.map(json_to_model, params=col, output={column: model})
 
     @classmethod
     def datasets(
@@ -443,9 +443,20 @@ class DataChain:
         )
         return listings(*args, **kwargs)
 
+    def persist(self) -> "Self":
+        """Saves temporary chain that will be removed after the process ends.
+        Temporary datasets are useful for optimization, for example when we have
+        multiple chains starting with identical sub-chain. We can then persist that
+        common chain and use it to calculate other chains, to avoid re-calculation
+        every time.
+        It returns the chain itself.
+        """
+        schema = self.signals_schema.clone_without_sys_signals().serialize()
+        return self._evolve(query=self._query.save(feature_schema=schema))
+
     def save(  # type: ignore[override]
         self,
-        name: Optional[str] = None,
+        name: str,
         version: Optional[int] = None,
         description: Optional[str] = None,
         labels: Optional[list[str]] = None,
@@ -454,8 +465,7 @@ class DataChain:
         """Save to a Dataset. It returns the chain itself.
 
         Parameters:
-            name : dataset name. Empty name saves to a temporary dataset that will be
-                removed after process ends. Temp dataset are useful for optimization.
+            name : dataset name.
             version : version of a dataset. Default - the last version that exist.
             description : description of a dataset.
             labels : labels of a dataset.
@@ -1112,7 +1122,7 @@ class DataChain:
         if self._query.attached:
             chain = self
         else:
-            chain = self.save()
+            chain = self.persist()
         assert chain.name is not None  # for mypy
         return PytorchDataset(
             chain.name,
@@ -1588,7 +1598,7 @@ class DataChain:
     def parse_tabular(
         self,
         output: OutputType = None,
-        object_name: str = "",
+        column: str = "",
         model_name: str = "",
         source: bool = True,
         nrows: Optional[int] = None,
@@ -1600,7 +1610,7 @@ class DataChain:
             output : Dictionary or feature class defining column names and their
                 corresponding types. List of column names is also accepted, in which
                 case types will be inferred.
-            object_name : Generated object column name.
+            column : Generated column name.
             model_name : Generated model name.
             source : Whether to include info about the source file.
             nrows : Optional row limit.
@@ -1651,14 +1661,14 @@ class DataChain:
                 raise DatasetPrepareError(self.name, e) from e
 
         if isinstance(output, dict):
-            model_name = model_name or object_name or ""
+            model_name = model_name or column or ""
             model = dict_to_data_model(model_name, output)
             output = model
         else:
             model = output  # type: ignore[assignment]
 
-        if object_name:
-            output = {object_name: model}  # type: ignore[dict-item]
+        if column:
+            output = {column: model}  # type: ignore[dict-item]
         elif isinstance(output, type(BaseModel)):
             output = {
                 name: info.annotation  # type: ignore[misc]