datachain 0.14.2__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/lib/dc/listings.py

@@ -1,25 +1,61 @@
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-)
+from typing import TYPE_CHECKING
 
+from datachain.lib.listing import LISTING_PREFIX, ls
 from datachain.lib.listing_info import ListingInfo
+from datachain.lib.settings import Settings
+from datachain.lib.signal_schema import SignalSchema
 from datachain.query import Session
+from datachain.query.dataset import DatasetQuery, QueryStep, step_result
 
 from .values import read_values
 
 if TYPE_CHECKING:
     from typing_extensions import ParamSpec
 
+    from datachain.dataset import DatasetVersion
+    from datachain.query.dataset import StepResult
+
     from .datachain import DataChain
 
     P = ParamSpec("P")
 
 
+class ReadOnlyQueryStep(QueryStep):
+    """
+    This step is used to read the dataset in read-only mode.
+    It is used to avoid the need to read the table metadata from the warehouse.
+    This is useful when we want to list the files in the dataset.
+    """
+
+    def apply(self) -> "StepResult":
+        import sqlalchemy as sa
+
+        def q(*columns):
+            return sa.select(*columns)
+
+        table_name = self.catalog.warehouse.dataset_table_name(
+            self.dataset, self.dataset_version
+        )
+        dataset_row_cls = self.catalog.warehouse.schema.dataset_row_cls
+        table = dataset_row_cls.new_table(
+            table_name,
+            columns=(
+                [
+                    *dataset_row_cls.sys_columns(),
+                    *dataset_row_cls.listing_columns(),
+                ]
+            ),
+        )
+
+        return step_result(
+            q, table.columns, dependencies=[(self.dataset, self.dataset_version)]
+        )
+
+
 def listings(
-    session: Optional[Session] = None,
+    session: Session | None = None,
     in_memory: bool = False,
-    object_name: str = "listing",
+    column: str = "listing",
     **kwargs,
 ) -> "DataChain":
     """Generate chain with list of cached listings.
@@ -38,6 +74,74 @@ def listings(
     return read_values(
         session=session,
         in_memory=in_memory,
-        output={object_name: ListingInfo},
-        **{object_name: catalog.listings()},  # type: ignore[arg-type]
+        output={column: ListingInfo},
+        **{column: catalog.listings()},  # type: ignore[arg-type]
     )
+
+
+def read_listing_dataset(
+    name: str,
+    version: str | None = None,
+    path: str = "",
+    session: Session | None = None,
+    settings: dict | None = None,
+) -> tuple["DataChain", "DatasetVersion"]:
+    """Read a listing dataset and return a DataChain and listing version.
+
+    Args:
+        name: Name of the dataset
+        version: Version of the dataset
+        path: Path within the listing to read. Path can have globs.
+        session: Optional Session object to use for reading
+        settings: Optional settings dictionary to use for reading
+
+    Returns:
+        tuple[DataChain, DatasetVersion]: A tuple containing:
+            - DataChain configured for listing files
+            - DatasetVersion object for the specified listing version
+
+    Example:
+        ```py
+        import datachain as dc
+        chain, listing_version = dc.read_listing_dataset(
+            "lst__s3://my-bucket/my-path", version="1.0.0", path="my-path"
+        )
+        chain.show()
+        ```
+    """
+    # Configure and return a DataChain for reading listing dataset files
+    # Uses ReadOnlyQueryStep to avoid warehouse metadata lookups
+    from datachain.lib.dc import Sys
+    from datachain.lib.file import File
+
+    from .datachain import DataChain
+
+    if not name.startswith(LISTING_PREFIX):
+        name = LISTING_PREFIX + name
+
+    session = Session.get(session)
+    dataset = session.catalog.get_dataset(name)
+    if version is None:
+        version = dataset.latest_version
+
+    query = DatasetQuery(name=name, session=session)
+
+    if settings:
+        cfg = {**settings}
+        if "prefetch" not in cfg:
+            cfg["prefetch"] = 0
+        _settings = Settings(**cfg)
+    else:
+        _settings = Settings(prefetch=0)
+    signal_schema = SignalSchema({"sys": Sys, "file": File})
+
+    query.starting_step = ReadOnlyQueryStep(query.catalog, dataset, version)
+    query.version = version
+    # We already know that this is a listing dataset,
+    # so we can set the listing function to True
+    query.set_listing_fn(lambda: True)
+
+    chain = DataChain(query, _settings, signal_schema)
+    chain = ls(chain, path, recursive=True, column="file")
+
+    return chain, dataset.get_version(version)

datachain/lib/dc/pandas.py

@@ -1,7 +1,4 @@
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-)
+from typing import TYPE_CHECKING
 
 from datachain.query import Session
 
@@ -19,10 +16,10 @@ if TYPE_CHECKING:
 def read_pandas(  # type: ignore[override]
     df: "pd.DataFrame",
     name: str = "",
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
-    object_name: str = "",
+    column: str = "",
 ) -> "DataChain":
     """Generate chain from pandas data-frame.
 
@@ -37,20 +34,27 @@ def read_pandas(  # type: ignore[override]
     """
     from .utils import DatasetPrepareError
 
-    fr_map = {col.lower(): df[col].tolist() for col in df.columns}
+    def get_col_name(col):
+        if isinstance(col, tuple):
+            # Join tuple elements with underscore for MultiIndex columns
+            return "_".join(map(str, col)).lower()
+        # Handle regular string column names
+        return str(col).lower()
 
-    for column in fr_map:
-        if not column.isidentifier():
+    fr_map = {get_col_name(col): df[col].tolist() for col in df.columns}
+
+    for c in fr_map:
+        if not c.isidentifier():
             raise DatasetPrepareError(
                 name,
-                f"import from pandas error - '{column}' cannot be a column name",
+                f"import from pandas error - '{c}' cannot be a column name",
             )
 
     return read_values(
         name,
         session,
         settings=settings,
-        object_name=object_name,
+        column=column,
         in_memory=in_memory,
         **fr_map,
     )

datachain/lib/dc/parquet.py

@@ -1,8 +1,5 @@
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Optional,
-)
+import os
+from typing import TYPE_CHECKING, Any
 
 from datachain.lib.data_model import DataType
 from datachain.query import Session
@@ -16,28 +13,34 @@ if TYPE_CHECKING:
 
 
 def read_parquet(
-    path,
+    path: str | os.PathLike[str] | list[str] | list[os.PathLike[str]],
     partitioning: Any = "hive",
-    output: Optional[dict[str, DataType]] = None,
-    object_name: str = "",
+    output: dict[str, DataType] | None = None,
+    column: str = "",
     model_name: str = "",
     source: bool = True,
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    session: Session | None = None,
+    settings: dict | None = None,
     **kwargs,
 ) -> "DataChain":
     """Generate chain from parquet files.
 
     Parameters:
-        path : Storage URI with directory. URI must start with storage prefix such
-            as `s3://`, `gs://`, `az://` or "file:///".
-        partitioning : Any pyarrow partitioning schema.
-        output : Dictionary defining column names and their corresponding types.
-        object_name : Created object column name.
-        model_name : Generated model name.
-        source : Whether to include info about the source file.
-        session : Session to use for the chain.
-        settings : Settings to use for the chain.
+        path: Storage path(s) or URI(s). Can be a local path or start with a
+            storage prefix like `s3://`, `gs://`, `az://`, `hf://` or "file:///".
+            Supports glob patterns:
+              - `*` : wildcard
+              - `**` : recursive wildcard
+              - `?` : single character
+              - `{a,b}` : brace expansion list
+              - `{1..9}` : brace numeric or alphabetic range
+        partitioning: Any pyarrow partitioning schema.
+        output: Dictionary defining column names and their corresponding types.
+        column: Created column name.
+        model_name: Generated model name.
+        source: Whether to include info about the source file.
+        session: Session to use for the chain.
+        settings: Settings to use for the chain.
 
     Example:
         Reading a single file:
@@ -46,10 +49,19 @@ def read_parquet(
         dc.read_parquet("s3://mybucket/file.parquet")
         ```
 
-        Reading a partitioned dataset from a directory:
+        All files from a directory:
         ```py
-        import datachain as dc
-        dc.read_parquet("s3://mybucket/dir")
+        dc.read_parquet("s3://mybucket/dir/")
+        ```
+
+        Only parquet files from a directory, and all it's subdirectories:
+        ```py
+        dc.read_parquet("s3://mybucket/dir/**/*.parquet")
+        ```
+
+        Using filename patterns - numeric, list, starting with zeros:
+        ```py
+        dc.read_parquet("s3://mybucket/202{1..4}/{yellow,green}-{01..12}.parquet")
         ```
     """
     from .storage import read_storage
@@ -57,7 +69,7 @@ def read_parquet(
     chain = read_storage(path, session=session, settings=settings, **kwargs)
     return chain.parse_tabular(
         output=output,
-        object_name=object_name,
+        column=column,
         model_name=model_name,
         source=source,
         format="parquet",

datachain/lib/dc/records.py

@@ -1,15 +1,10 @@
-from typing import (
-    TYPE_CHECKING,
-    Optional,
-    Union,
-)
+from collections.abc import Iterable
+from typing import TYPE_CHECKING
 
 import sqlalchemy
 
 from datachain.lib.data_model import DataType
-from datachain.lib.file import (
-    File,
-)
+from datachain.lib.file import File
 from datachain.lib.signal_schema import SignalSchema
 from datachain.query import Session
 
@@ -20,29 +15,37 @@ if TYPE_CHECKING:
 
     P = ParamSpec("P")
 
+READ_RECORDS_BATCH_SIZE = 10000
+
 
 def read_records(
-    to_insert: Optional[Union[dict, list[dict]]],
-    session: Optional[Session] = None,
-    settings: Optional[dict] = None,
+    to_insert: dict | Iterable[dict] | None,
+    session: Session | None = None,
+    settings: dict | None = None,
     in_memory: bool = False,
-    schema: Optional[dict[str, DataType]] = None,
+    schema: dict[str, DataType] | None = None,
 ) -> "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
     or other sources.
 
     Parameters:
-        to_insert : records (or a single record) to insert. Each record is
-                    a dictionary of signals and theirs values.
-        schema : describes chain signals and their corresponding types
+        to_insert: records (or a single record) to insert. Each record is
+                    a dictionary of signals and their values.
+        schema: describes chain signals and their corresponding types
 
     Example:
         ```py
         import datachain as dc
         single_record = dc.read_records(dc.DEFAULT_FILE_RECORD)
         ```
+
+    Notes:
+        This call blocks until all records are inserted.
     """
+    from datachain.query.dataset import adjust_outputs, get_col_types
+    from datachain.sql.types import SQLType
+
     from .datasets import read_dataset
 
     session = Session.get(session, in_memory=in_memory)
@@ -56,7 +59,7 @@ def read_records(
         signal_schema = SignalSchema(schema)
         columns = [
             sqlalchemy.Column(c.name, c.type)  # type: ignore[union-attr]
-            for c in signal_schema.db_signals(as_columns=True)  # type: ignore[assignment]
+            for c in signal_schema.db_signals(as_columns=True)
         ]
     else:
         columns = [
@@ -66,6 +69,7 @@ def read_records(
 
     dsr = catalog.create_dataset(
         name,
+        catalog.metastore.default_project,
         columns=columns,
         feature_schema=(
             signal_schema.clone_without_sys_signals().serialize()
@@ -74,8 +78,6 @@ def read_records(
         ),
     )
 
-    session.add_dataset_version(dsr, dsr.latest_version)
-
     if isinstance(to_insert, dict):
         to_insert = [to_insert]
     elif not to_insert:
@@ -83,8 +85,14 @@ def read_records(
 
     warehouse = catalog.warehouse
     dr = warehouse.dataset_rows(dsr)
-    db = warehouse.db
-    insert_q = dr.get_table().insert()
-    for record in to_insert:
-        db.execute(insert_q.values(**record))
-    return read_dataset(name=dsr.name, session=session, settings=settings)
+    table = dr.get_table()
+
+    # Optimization: Compute row types once, rather than for every row.
+    col_types = get_col_types(
+        warehouse,
+        {c.name: c.type for c in columns if isinstance(c.type, SQLType)},
+    )
+    records = (adjust_outputs(warehouse, record, col_types) for record in to_insert)
+    warehouse.insert_rows(table, records, batch_size=READ_RECORDS_BATCH_SIZE)
+    warehouse.insert_rows_done(table)
+    return read_dataset(name=dsr.full_name, session=session, settings=settings)