datachain 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

datachain/lib/data_model.py

@@ -18,6 +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]
 
 
 class DataModel(BaseModel):

datachain/lib/dc.py

@@ -309,6 +309,7 @@ class DataChain(DatasetQuery):
         *,
         type: Literal["binary", "text", "image"] = "binary",
         session: Optional[Session] = None,
+        in_memory: bool = False,
         recursive: Optional[bool] = True,
         object_name: str = "file",
         update: bool = False,
@@ -332,7 +333,14 @@ class DataChain(DatasetQuery):
         """
         func = get_file(type)
         return (
-            cls(path, session=session, recursive=recursive, update=update, **kwargs)
+            cls(
+                path,
+                session=session,
+                recursive=recursive,
+                update=update,
+                in_memory=in_memory,
+                **kwargs,
+            )
             .map(**{object_name: func})
             .select(object_name)
         )
@@ -479,7 +487,10 @@ class DataChain(DatasetQuery):
 
     @classmethod
     def datasets(
-        cls, session: Optional[Session] = None, object_name: str = "dataset"
+        cls,
+        session: Optional[Session] = None,
+        in_memory: bool = False,
+        object_name: str = "dataset",
     ) -> "DataChain":
         """Generate chain with list of registered datasets.
 
@@ -492,7 +503,7 @@ class DataChain(DatasetQuery):
                 print(f"{ds.name}@v{ds.version}")
             ```
         """
-        session = Session.get(session)
+        session = Session.get(session, in_memory=in_memory)
         catalog = session.catalog
 
         datasets = [
@@ -502,13 +513,14 @@ class DataChain(DatasetQuery):
 
         return cls.from_values(
             session=session,
+            in_memory=in_memory,
             output={object_name: DatasetInfo},
             **{object_name: datasets},  # type: ignore[arg-type]
         )
 
     def print_json_schema(  # type: ignore[override]
         self, jmespath: Optional[str] = None, model_name: Optional[str] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Print JSON data model and save it. It returns the chain itself.
 
         Parameters:
@@ -533,7 +545,7 @@ class DataChain(DatasetQuery):
 
     def print_jsonl_schema(  # type: ignore[override]
         self, jmespath: Optional[str] = None, model_name: Optional[str] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Print JSON data model and save it. It returns the chain itself.
 
         Parameters:
@@ -549,7 +561,7 @@ class DataChain(DatasetQuery):
 
     def save(  # type: ignore[override]
         self, name: Optional[str] = None, version: Optional[int] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
 
         Parameters:
@@ -785,7 +797,7 @@ class DataChain(DatasetQuery):
             descending (bool): Whether to sort in descending order or not.
         """
         if descending:
-            args = tuple([sqlalchemy.desc(a) for a in args])
+            args = tuple(sqlalchemy.desc(a) for a in args)
 
         return super().order_by(*args)
 
@@ -1142,6 +1154,7 @@ class DataChain(DatasetQuery):
         cls,
         ds_name: str = "",
         session: Optional[Session] = None,
+        in_memory: bool = False,
         output: OutputType = None,
         object_name: str = "",
         **fr_map,
@@ -1158,7 +1171,9 @@ class DataChain(DatasetQuery):
         def _func_fr() -> Iterator[tuple_type]:  # type: ignore[valid-type]
             yield from tuples
 
-        chain = DataChain.from_records(DataChain.DEFAULT_FILE_RECORD, session=session)
+        chain = DataChain.from_records(
+            DataChain.DEFAULT_FILE_RECORD, session=session, in_memory=in_memory
+        )
         if object_name:
             output = {object_name: DataChain._dict_to_data_model(object_name, output)}  # type: ignore[arg-type]
         return chain.gen(_func_fr, output=output)
@@ -1169,6 +1184,7 @@ class DataChain(DatasetQuery):
         df: "pd.DataFrame",
         name: str = "",
         session: Optional[Session] = None,
+        in_memory: bool = False,
         object_name: str = "",
     ) -> "DataChain":
         """Generate chain from pandas data-frame.
@@ -1196,7 +1212,9 @@ class DataChain(DatasetQuery):
                     f"import from pandas error - '{column}' cannot be a column name",
                 )
 
-        return cls.from_values(name, session, object_name=object_name, **fr_map)
+        return cls.from_values(
+            name, session, object_name=object_name, in_memory=in_memory, **fr_map
+        )
 
     def to_pandas(self, flatten=False) -> "pd.DataFrame":
         """Return a pandas DataFrame from the chain.
@@ -1206,14 +1224,14 @@ class DataChain(DatasetQuery):
         """
         headers, max_length = self._effective_signals_schema.get_headers_with_length()
         if flatten or max_length < 2:
-            df = pd.DataFrame.from_records(self.to_records())
+            columns = []
             if headers:
-                df.columns = [".".join(filter(None, header)) for header in headers]
-            return df
+                columns = [".".join(filter(None, header)) for header in headers]
+            return pd.DataFrame.from_records(self.to_records(), columns=columns)
 
-        transposed_result = list(map(list, zip(*self.results())))
-        data = {tuple(n): val for n, val in zip(headers, transposed_result)}
-        return pd.DataFrame(data)
+        return pd.DataFrame(
+            self.results(), columns=pd.MultiIndex.from_tuples(map(tuple, headers))
+        )
 
     def show(
         self,
@@ -1232,6 +1250,12 @@ class DataChain(DatasetQuery):
         """
         dc = self.limit(limit) if limit > 0 else self
         df = dc.to_pandas(flatten)
+
+        if df.empty:
+            print("Empty result")
+            print(f"Columns: {list(df.columns)}")
+            return
+
         if transpose:
             df = df.T
 
@@ -1270,7 +1294,7 @@ class DataChain(DatasetQuery):
         source: bool = True,
         nrows: Optional[int] = None,
         **kwargs,
-    ) -> "DataChain":
+    ) -> "Self":
         """Generate chain from list of tabular files.
 
         Parameters:
@@ -1390,7 +1414,8 @@ class DataChain(DatasetQuery):
             dc = DataChain.from_csv("s3://mybucket/dir")
             ```
         """
-        from pyarrow.csv import ParseOptions, ReadOptions
+        from pandas.io.parsers.readers import STR_NA_VALUES
+        from pyarrow.csv import ConvertOptions, ParseOptions, ReadOptions
         from pyarrow.dataset import CsvFileFormat
 
         chain = DataChain.from_storage(path, **kwargs)
@@ -1414,7 +1439,14 @@ class DataChain(DatasetQuery):
 
         parse_options = ParseOptions(delimiter=delimiter)
         read_options = ReadOptions(column_names=column_names)
-        format = CsvFileFormat(parse_options=parse_options, read_options=read_options)
+        convert_options = ConvertOptions(
+            strings_can_be_null=True, null_values=STR_NA_VALUES
+        )
+        format = CsvFileFormat(
+            parse_options=parse_options,
+            read_options=read_options,
+            convert_options=convert_options,
+        )
         return chain.parse_tabular(
             output=output,
             object_name=object_name,
@@ -1491,6 +1523,7 @@ class DataChain(DatasetQuery):
         cls,
         to_insert: Optional[Union[dict, list[dict]]],
         session: Optional[Session] = None,
+        in_memory: bool = False,
     ) -> "DataChain":
         """Create a DataChain from the provided records. This method can be used for
         programmatically generating a chain in contrast of reading data from storages
@@ -1506,7 +1539,7 @@ class DataChain(DatasetQuery):
             single_record = DataChain.from_records(DataChain.DEFAULT_FILE_RECORD)
             ```
         """
-        session = Session.get(session)
+        session = Session.get(session, in_memory=in_memory)
         catalog = session.catalog
 
         name = session.generate_temp_dataset_name()

datachain/lib/listing.py

@@ -0,0 +1,111 @@
+import asyncio
+from collections.abc import AsyncIterator, Iterator, Sequence
+from typing import Callable, Optional
+
+from botocore.exceptions import ClientError
+from fsspec.asyn import get_loop
+
+from datachain.asyn import iter_over_async
+from datachain.client import Client
+from datachain.error import ClientError as DataChainClientError
+from datachain.lib.file import File
+
+ResultQueue = asyncio.Queue[Optional[Sequence[File]]]
+
+DELIMITER = "/"  # Path delimiter
+FETCH_WORKERS = 100
+
+
+async def _fetch_dir(client, prefix, result_queue) -> set[str]:
+    path = f"{client.name}/{prefix}"
+    infos = await client.ls_dir(path)
+    files = []
+    subdirs = set()
+    for info in infos:
+        full_path = info["name"]
+        subprefix = client.rel_path(full_path)
+        if prefix.strip(DELIMITER) == subprefix.strip(DELIMITER):
+            continue
+        if info["type"] == "directory":
+            subdirs.add(subprefix)
+        else:
+            files.append(client.info_to_file(info, subprefix))
+    if files:
+        await result_queue.put(files)
+    return subdirs
+
+
+async def _fetch(
+    client, start_prefix: str, result_queue: ResultQueue, fetch_workers
+) -> None:
+    loop = get_loop()
+
+    queue: asyncio.Queue[str] = asyncio.Queue()
+    queue.put_nowait(start_prefix)
+
+    async def process(queue) -> None:
+        while True:
+            prefix = await queue.get()
+            try:
+                subdirs = await _fetch_dir(client, prefix, result_queue)
+                for subdir in subdirs:
+                    queue.put_nowait(subdir)
+            except Exception:
+                while not queue.empty():
+                    queue.get_nowait()
+                    queue.task_done()
+                raise
+
+            finally:
+                queue.task_done()
+
+    try:
+        workers: list[asyncio.Task] = [
+            loop.create_task(process(queue)) for _ in range(fetch_workers)
+        ]
+
+        # Wait for all fetch tasks to complete
+        await queue.join()
+        # Stop the workers
+        excs = []
+        for worker in workers:
+            if worker.done() and (exc := worker.exception()):
+                excs.append(exc)
+            else:
+                worker.cancel()
+        if excs:
+            raise excs[0]
+    except ClientError as exc:
+        raise DataChainClientError(
+            exc.response.get("Error", {}).get("Message") or exc,
+            exc.response.get("Error", {}).get("Code"),
+        ) from exc
+    finally:
+        # This ensures the progress bar is closed before any exceptions are raised
+        result_queue.put_nowait(None)
+
+
+async def _scandir(client, prefix, fetch_workers) -> AsyncIterator:
+    """Recursively goes through dir tree and yields files"""
+    result_queue: ResultQueue = asyncio.Queue()
+    loop = get_loop()
+    main_task = loop.create_task(_fetch(client, prefix, result_queue, fetch_workers))
+    while (files := await result_queue.get()) is not None:
+        for f in files:
+            yield f
+
+    await main_task
+
+
+def list_bucket(uri: str, client_config=None, fetch_workers=FETCH_WORKERS) -> Callable:
+    """
+    Function that returns another generator function that yields File objects
+    from bucket where each File represents one bucket entry.
+    """
+
+    def list_func() -> Iterator[File]:
+        config = client_config or {}
+        client, path = Client.parse_url(uri, None, **config)  # type: ignore[arg-type]
+        yield from iter_over_async(_scandir(client, path, fetch_workers), get_loop())
+
+    return list_func

datachain/lib/meta_formats.py

@@ -11,12 +11,16 @@ from collections.abc import Iterator
 from typing import Any, Callable
 
 import jmespath as jsp
-from pydantic import Field, ValidationError  # noqa: F401
+from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401
 
 from datachain.lib.data_model import DataModel  # noqa: F401
 from datachain.lib.file import File
 
 
+class UserModel(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+
 def generate_uuid():
     return uuid.uuid4()  # Generates a random UUID.
 
@@ -72,6 +76,8 @@ def read_schema(source_file, data_type="csv", expr=None, model_name=None):
         data_type,
         "--class-name",
         model_name,
+        "--base-class",
+        "datachain.lib.meta_formats.UserModel",
     ]
     try:
         result = subprocess.run(  # noqa: S603
@@ -87,7 +93,7 @@ def read_schema(source_file, data_type="csv", expr=None, model_name=None):
     except subprocess.CalledProcessError as e:
         model_output = f"An error occurred in datamodel-codegen: {e.stderr}"
     print(f"{model_output}")
-    print("\n" + "from datachain.lib.data_model import DataModel" + "\n")
+    print("from datachain.lib.data_model import DataModel")
     print("\n" + f"DataModel.register({model_name})" + "\n")
     print("\n" + f"spec={model_name}" + "\n")
     return model_output

datachain/node.py

@@ -47,7 +47,7 @@ class DirTypeGroup:
 @attrs.define
 class Node:
     sys__id: int = 0
-    sys__rand: int = -1
+    sys__rand: int = 0
     vtype: str = ""
     dir_type: Optional[int] = None
     path: str = ""

datachain/query/dataset.py

@@ -34,6 +34,7 @@ from sqlalchemy.sql.elements import ColumnClause, ColumnElement
 from sqlalchemy.sql.expression import label
 from sqlalchemy.sql.schema import TableClause
 from sqlalchemy.sql.selectable import Select
+from tqdm import tqdm
 
 from datachain.asyn import ASYNC_WORKERS, AsyncMapper, OrderedMapper
 from datachain.catalog import (
@@ -125,7 +126,10 @@ class QueryGenerator:
     func: QueryGeneratorFunc
     columns: tuple[ColumnElement, ...]
 
-    def exclude(self, column_names) -> Select:
+    def only(self, column_names: Sequence[str]) -> Select:
+        return self.func(*(c for c in self.columns if c.name in column_names))
+
+    def exclude(self, column_names: Sequence[str]) -> Select:
         return self.func(*(c for c in self.columns if c.name not in column_names))
 
     def select(self, column_names=None) -> Select:
@@ -465,6 +469,12 @@ class UDFStep(Step, ABC):
 
         try:
             if workers:
+                if self.catalog.in_memory:
+                    raise RuntimeError(
+                        "In-memory databases cannot be used with "
+                        "distributed processing."
+                    )
+
                 from datachain.catalog.loader import get_distributed_class
 
                 distributor = get_distributed_class(min_task_size=self.min_task_size)
@@ -482,6 +492,10 @@ class UDFStep(Step, ABC):
                 )
             elif processes:
                 # Parallel processing (faster for more CPU-heavy UDFs)
+                if self.catalog.in_memory:
+                    raise RuntimeError(
+                        "In-memory databases cannot be used with parallel processing."
+                    )
                 udf_info = {
                     "udf_data": filtered_cloudpickle_dumps(self.udf),
                     "catalog_init": self.catalog.get_init_params(),
@@ -1049,6 +1063,7 @@ class DatasetQuery:
         indexing_feature_schema: Optional[dict] = None,
         indexing_column_types: Optional[dict[str, Any]] = None,
         update: Optional[bool] = False,
+        in_memory: bool = False,
     ):
         if client_config is None:
             client_config = {}
@@ -1057,7 +1072,7 @@ class DatasetQuery:
             client_config["anon"] = True
 
         self.session = Session.get(
-            session, catalog=catalog, client_config=client_config
+            session, catalog=catalog, client_config=client_config, in_memory=in_memory
         )
         self.catalog = catalog or self.session.catalog
         self.steps: list[Step] = []
@@ -1648,18 +1663,13 @@ class DatasetQuery:
 
             dr = self.catalog.warehouse.dataset_rows(dataset)
 
-            # Exclude the id column and let the db create it to avoid unique
-            # constraint violations.
-            q = query.exclude(("sys__id",))
-            if q._order_by_clauses:
-                # ensuring we have id sorted by order by clause if it exists in a query
-                q = q.add_columns(
-                    f.row_number().over(order_by=q._order_by_clauses).label("sys__id")
+            with tqdm(desc="Saving", unit=" rows") as pbar:
+                self.catalog.warehouse.copy_table(
+                    dr.get_table(),
+                    query.select(),
+                    progress_cb=pbar.update,
                 )
 
-            cols = tuple(c.name for c in q.selected_columns)
-            insert_q = sqlalchemy.insert(dr.get_table()).from_select(cols, q)
-            self.catalog.warehouse.db.execute(insert_q, **kwargs)
             self.catalog.metastore.update_dataset_status(
                 dataset, DatasetStatus.COMPLETE, version=version
             )

datachain/query/schema.py

@@ -45,6 +45,10 @@ class Column(sa.ColumnClause, metaclass=ColumnMeta):
         """Search for matches using glob pattern matching."""
         return self.op("GLOB")(glob_str)
 
+    def regexp(self, regexp_str):
+        """Search for matches using regexp pattern matching."""
+        return self.op("REGEXP")(regexp_str)
+
 
 class UDFParameter(ABC):
     @abstractmethod

datachain/query/session.py

@@ -46,6 +46,7 @@ class Session:
         name="",
         catalog: Optional["Catalog"] = None,
         client_config: Optional[dict] = None,
+        in_memory: bool = False,
     ):
         if re.match(r"^[0-9a-zA-Z]+$", name) is None:
             raise ValueError(
@@ -58,7 +59,9 @@ class Session:
         session_uuid = uuid4().hex[: self.SESSION_UUID_LEN]
         self.name = f"{name}_{session_uuid}"
         self.is_new_catalog = not catalog
-        self.catalog = catalog or get_catalog(client_config=client_config)
+        self.catalog = catalog or get_catalog(
+            client_config=client_config, in_memory=in_memory
+        )
 
     def __enter__(self):
         return self
@@ -89,6 +92,7 @@ class Session:
         session: Optional["Session"] = None,
         catalog: Optional["Catalog"] = None,
         client_config: Optional[dict] = None,
+        in_memory: bool = False,
     ) -> "Session":
         """Creates a Session() object from a catalog.
 
@@ -102,7 +106,10 @@ class Session:
 
         if cls.GLOBAL_SESSION is None:
             cls.GLOBAL_SESSION_CTX = Session(
-                cls.GLOBAL_SESSION_NAME, catalog, client_config=client_config
+                cls.GLOBAL_SESSION_NAME,
+                catalog,
+                client_config=client_config,
+                in_memory=in_memory,
             )
             cls.GLOBAL_SESSION = cls.GLOBAL_SESSION_CTX.__enter__()
             atexit.register(cls._global_cleanup)

datachain/sql/default/base.py

@@ -1,8 +1,10 @@
 from datachain.sql.types import (
+    DBDefaults,
     TypeConverter,
     TypeDefaults,
     TypeReadConverter,
     register_backend_types,
+    register_db_defaults,
     register_type_defaults,
     register_type_read_converters,
 )
@@ -18,5 +20,6 @@ def setup() -> None:
     register_backend_types("default", TypeConverter())
     register_type_read_converters("default", TypeReadConverter())
     register_type_defaults("default", TypeDefaults())
+    register_db_defaults("default", DBDefaults())
 
     setup_is_complete = True

datachain/sql/sqlite/base.py

@@ -22,8 +22,10 @@ from datachain.sql.sqlite.types import (
     register_type_converters,
 )
 from datachain.sql.types import (
+    DBDefaults,
     TypeDefaults,
     register_backend_types,
+    register_db_defaults,
     register_type_defaults,
     register_type_read_converters,
 )
@@ -66,6 +68,7 @@ def setup():
     register_backend_types("sqlite", SQLiteTypeConverter())
     register_type_read_converters("sqlite", SQLiteTypeReadConverter())
     register_type_defaults("sqlite", TypeDefaults())
+    register_db_defaults("sqlite", DBDefaults())
 
     compiles(sql_path.parent, "sqlite")(compile_path_parent)
     compiles(sql_path.name, "sqlite")(compile_path_name)
@@ -218,19 +221,45 @@ def path_name(path):
     return func.ltrim(func.substr(path, func.length(path_parent(path)) + 1), slash)
 
 
-def path_file_ext_length(path):
-    name = path_name(path)
+def name_file_ext_length(name):
     expr = func.length(name) - func.length(
         func.rtrim(name, func.replace(name, dot, empty_str))
     )
     return case((func.instr(name, dot) == 0, 0), else_=expr)
 
 
+def path_file_ext_length(path):
+    name = path_name(path)
+    return name_file_ext_length(name)
+
+
 def path_file_stem(path):
-    return func.rtrim(
-        func.substr(path, 1, func.length(path) - path_file_ext_length(path)), dot
+    path_length = func.length(path)
+    parent_length = func.length(path_parent(path))
+
+    name_expr = func.rtrim(
+        func.substr(
+            path,
+            1,
+            path_length - name_file_ext_length(path),
+        ),
+        dot,
+    )
+
+    full_path_expr = func.ltrim(
+        func.rtrim(
+            func.substr(
+                path,
+                parent_length + 1,
+                path_length - parent_length - path_file_ext_length(path),
+            ),
+            dot,
+        ),
+        slash,
     )
 
+    return case((func.instr(path, slash) == 0, name_expr), else_=full_path_expr)
+
 
 def path_file_ext(path):
     return func.substr(path, func.length(path) - path_file_ext_length(path) + 1)