datachain 0.27.0__tar.gz → 0.28.1__tar.gz

{datachain-0.27.0 → datachain-0.28.1}/.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.12.4'
+    rev: 'v0.12.5'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.27.0 → datachain-0.28.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.27.0
+Version: 0.28.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -45,7 +45,7 @@ Requires-Dist: datamodel-code-generator>=0.25
 Requires-Dist: Pillow<12,>=10.0.0
 Requires-Dist: msgpack<2,>=1.0.4
 Requires-Dist: psutil
-Requires-Dist: huggingface_hub
+Requires-Dist: huggingface_hub<0.34.0
 Requires-Dist: iterative-telemetry>=0.0.10
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21

{datachain-0.27.0 → datachain-0.28.1}/docs/commands/job/run.md

@@ -14,6 +14,7 @@ usage: datachain job run [-h] [-v] [-q] [--team TEAM] [--env-file ENV_FILE]
                          [--req-file REQ_FILE] [--req REQ [REQ ...]]
                          [--priority PRIORITY]
                          [--start-time START_TIME] [--cron CRON]
+                         [--no-wait]
                          file
 ```
 
@@ -40,6 +41,7 @@ This command runs a job in Studio using the specified query file. You can config
 * `--priority PRIORITY` - Priority for the job in range 0-5. Lower value is higher priority (default: 5)
 * `--start-time START_TIME` - Time to schedule the task in YYYY-MM-DDTHH:mm format or natural language.
 * `--cron CRON` - Cron expression for the cron task.
+* `--no-wait` - Do not wait for the job to finish.
 * `-h`, `--help` - Show the help message and exit.
 * `-v`, `--verbose` - Be verbose.
 * `-q`, `--quiet` - Be quiet.
@@ -131,6 +133,12 @@ datachain job run --cron "@monthly" query.py
 datachain job run --start-time "tomorrow 3pm" --cron "0 0 * * *" query.py
 ```
 
+12. Start the job and do not wait for the job to complete
+```bash
+# Do not follow or tail the logs from Studio.
+datachain job run query.py --no-wait
+```
+
 ## Notes
 
 * Closing the logs command (e.g., with Ctrl+C) will only stop displaying the logs but will not cancel the job execution

{datachain-0.27.0 → datachain-0.28.1}/mkdocs.yml

@@ -177,7 +177,7 @@ plugins:
             - https://numpy.org/doc/stable/objects.inv
             - https://pandas.pydata.org/docs/objects.inv
             - https://arrow.apache.org/docs/objects.inv
-            - https://docs.sqlalchemy.org/objects.inv
+            # - https://docs.sqlalchemy.org/objects.inv  # SSL certificate issue
             - https://docs.pydantic.dev/latest/objects.inv
 
 watch:

{datachain-0.27.0 → datachain-0.28.1}/pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
   "Pillow>=10.0.0,<12",
   "msgpack>=1.0.4,<2",
   "psutil",
-  "huggingface_hub",
+  "huggingface_hub<0.34.0",
   "iterative-telemetry>=0.0.10",
   "platformdirs",
   "dvc-studio-client>=0.21,<1",

{datachain-0.27.0 → datachain-0.28.1}/src/datachain/cli/parser/job.py

@@ -109,6 +109,11 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
     studio_run_parser.add_argument(
         "--cron", action="store", help="Cron expression for the cron task."
     )
+    studio_run_parser.add_argument(
+        "--no-wait",
+        action="store_true",
+        help="Do not wait for the job to finish",
+    )
 
     studio_ls_help = "List jobs in Studio"
     studio_ls_description = "List jobs in Studio."

{datachain-0.27.0 → datachain-0.28.1}/src/datachain/data_storage/job.py

@@ -12,10 +12,11 @@ class JobStatus(int, Enum):
     CANCELING = 7
     CANCELED = 8
     CANCELING_SCHEDULED = 9
+    TASK = 11
 
     @classmethod
     def finished(cls) -> tuple[int, ...]:
-        return cls.COMPLETE, cls.FAILED, cls.CANCELED
+        return cls.COMPLETE, cls.FAILED, cls.CANCELED, cls.TASK
 
 
 class JobQueryType(int, Enum):

{datachain-0.27.0 → datachain-0.28.1}/src/datachain/lib/dc/datachain.py

@@ -2419,9 +2419,11 @@ class DataChain:
             ds.to_storage("gs://mybucket", placement="filename")
             ```
         """
+        chain = self.persist()
+        count = chain.count()
+
         if placement == "filename" and (
-            self._query.distinct(pathfunc.name(C(f"{signal}__path"))).count()
-            != self._query.count()
+            chain._query.distinct(pathfunc.name(C(f"{signal}__path"))).count() != count
         ):
             raise ValueError("Files with the same name found")
 
@@ -2433,7 +2435,7 @@ class DataChain:
             unit=" files",
             unit_scale=True,
             unit_divisor=10,
-            total=self.count(),
+            total=count,
             leave=False,
         )
         file_exporter = FileExporter(
@@ -2444,7 +2446,10 @@ class DataChain:
             max_threads=num_threads or 1,
             client_config=client_config,
         )
-        file_exporter.run(self.to_values(signal), progress_bar)
+        file_exporter.run(
+            (rows[0] for rows in chain.to_iter(signal)),
+            progress_bar,
+        )
 
     def shuffle(self) -> "Self":
         """Shuffle the rows of the chain deterministically."""

{datachain-0.27.0 → datachain-0.28.1}/src/datachain/lib/file.py

@@ -23,7 +23,7 @@ from pydantic import Field, field_validator
 
 from datachain.client.fileslice import FileSlice
 from datachain.lib.data_model import DataModel
-from datachain.lib.utils import DataChainError
+from datachain.lib.utils import DataChainError, rebase_path
 from datachain.nodes_thread_pool import NodesThreadPool
 from datachain.sql.types import JSON, Boolean, DateTime, Int, String
 from datachain.utils import TIME_ZERO
@@ -634,6 +634,40 @@ class File(DataModel):
             location=self.location,
         )
 
+    def rebase(
+        self,
+        old_base: str,
+        new_base: str,
+        suffix: str = "",
+        extension: str = "",
+    ) -> str:
+        """
+        Rebase the file's URI from one base directory to another.
+
+        Args:
+            old_base: Base directory to remove from the file's URI
+            new_base: New base directory to prepend
+            suffix: Optional suffix to add before file extension
+            extension: Optional new file extension (without dot)
+
+        Returns:
+            str: Rebased URI with new base directory
+
+        Raises:
+            ValueError: If old_base is not found in the file's URI
+
+        Examples:
+            >>> file = File(source="s3://bucket", path="data/2025-05-27/file.wav")
+            >>> file.rebase("s3://bucket/data", "s3://output-bucket/processed", \
+                    extension="mp3")
+            's3://output-bucket/processed/2025-05-27/file.mp3'
+
+            >>> file.rebase("data/audio", "/local/output", suffix="_ch1",
+                    extension="npy")
+            '/local/output/file_ch1.npy'
+        """
+        return rebase_path(self.get_uri(), old_base, new_base, suffix, extension)
+
 
 def resolve(file: File) -> File:
     """
@@ -1219,6 +1253,24 @@ class Audio(DataModel):
     codec: str = Field(default="")
     bit_rate: int = Field(default=-1)
 
+    @staticmethod
+    def get_channel_name(num_channels: int, channel_idx: int) -> str:
+        """Map channel index to meaningful name based on common audio formats"""
+        channel_mappings = {
+            1: ["Mono"],
+            2: ["Left", "Right"],
+            4: ["W", "X", "Y", "Z"],  # First-order Ambisonics
+            6: ["FL", "FR", "FC", "LFE", "BL", "BR"],  # 5.1 surround
+            8: ["FL", "FR", "FC", "LFE", "BL", "BR", "SL", "SR"],  # 7.1 surround
+        }
+
+        if num_channels in channel_mappings:
+            channels = channel_mappings[num_channels]
+            if 0 <= channel_idx < len(channels):
+                return channels[channel_idx]
+
+        return f"Ch{channel_idx + 1}"
+
 
 class ArrowRow(DataModel):
     """`DataModel` for reading row from Arrow-supported file."""

datachain-0.28.1/src/datachain/lib/utils.py

@@ -0,0 +1,155 @@
+import re
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from pathlib import PurePosixPath
+from urllib.parse import urlparse
+
+
+class AbstractUDF(ABC):
+    @abstractmethod
+    def process(self, *args, **kwargs):
+        pass
+
+    @abstractmethod
+    def setup(self):
+        pass
+
+    @abstractmethod
+    def teardown(self):
+        pass
+
+
+class DataChainError(Exception):
+    pass
+
+
+class DataChainParamsError(DataChainError):
+    pass
+
+
+class DataChainColumnError(DataChainParamsError):
+    def __init__(self, col_name: str, msg: str):
+        super().__init__(f"Error for column {col_name}: {msg}")
+
+
+def normalize_col_names(col_names: Sequence[str]) -> dict[str, str]:
+    """Returns normalized_name -> original_name dict."""
+    gen_col_counter = 0
+    new_col_names = {}
+    org_col_names = set(col_names)
+
+    for org_column in col_names:
+        new_column = org_column.lower()
+        new_column = re.sub("[^0-9a-z]+", "_", new_column)
+        new_column = new_column.strip("_")
+
+        generated_column = new_column
+
+        while (
+            not generated_column.isidentifier()
+            or generated_column in new_col_names
+            or (generated_column != org_column and generated_column in org_col_names)
+        ):
+            if new_column:
+                generated_column = f"c{gen_col_counter}_{new_column}"
+            else:
+                generated_column = f"c{gen_col_counter}"
+            gen_col_counter += 1
+
+        new_col_names[generated_column] = org_column
+
+    return new_col_names
+
+
+def rebase_path(
+    src_path: str,
+    old_base: str,
+    new_base: str,
+    suffix: str = "",
+    extension: str = "",
+) -> str:
+    """
+    Rebase a file path from one base directory to another.
+
+    Args:
+        src_path: Source file path (can include URI scheme like s3://)
+        old_base: Base directory to remove from src_path
+        new_base: New base directory to prepend
+        suffix: Optional suffix to add before file extension
+        extension: Optional new file extension (without dot)
+
+    Returns:
+        str: Rebased path with new base directory
+
+    Raises:
+        ValueError: If old_base is not found in src_path
+    """
+    # Parse URIs to handle schemes properly
+    src_parsed = urlparse(src_path)
+    old_base_parsed = urlparse(old_base)
+    new_base_parsed = urlparse(new_base)
+
+    # Get the path component (without scheme)
+    if src_parsed.scheme:
+        src_path_only = src_parsed.netloc + src_parsed.path
+    else:
+        src_path_only = src_path
+
+    if old_base_parsed.scheme:
+        old_base_only = old_base_parsed.netloc + old_base_parsed.path
+    else:
+        old_base_only = old_base
+
+    # Normalize paths
+    src_path_norm = PurePosixPath(src_path_only).as_posix()
+    old_base_norm = PurePosixPath(old_base_only).as_posix()
+
+    # Find where old_base appears in src_path
+    if old_base_norm in src_path_norm:
+        # Find the index where old_base appears
+        idx = src_path_norm.find(old_base_norm)
+        if idx == -1:
+            raise ValueError(f"old_base '{old_base}' not found in src_path")
+
+        # Extract the relative path after old_base
+        relative_start = idx + len(old_base_norm)
+        # Skip leading slash if present
+        if relative_start < len(src_path_norm) and src_path_norm[relative_start] == "/":
+            relative_start += 1
+        relative_path = src_path_norm[relative_start:]
+    else:
+        raise ValueError(f"old_base '{old_base}' not found in src_path")
+
+    # Parse the filename
+    path_obj = PurePosixPath(relative_path)
+    stem = path_obj.stem
+    current_ext = path_obj.suffix
+
+    # Apply suffix and extension changes
+    new_stem = stem + suffix if suffix else stem
+    if extension:
+        new_ext = f".{extension}"
+    elif current_ext:
+        new_ext = current_ext
+    else:
+        new_ext = ""
+
+    # Build new filename
+    new_name = new_stem + new_ext
+
+    # Reconstruct path with new base
+    parent = str(path_obj.parent)
+    if parent == ".":
+        new_relative_path = new_name
+    else:
+        new_relative_path = str(PurePosixPath(parent) / new_name)
+
+    # Handle new_base URI scheme
+    if new_base_parsed.scheme:
+        # Has schema like s3://
+        base_path = new_base_parsed.netloc + new_base_parsed.path
+        base_path = PurePosixPath(base_path).as_posix()
+        full_path = str(PurePosixPath(base_path) / new_relative_path)
+        return f"{new_base_parsed.scheme}://{full_path}"
+    # Regular path
+    return str(PurePosixPath(new_base) / new_relative_path)

{datachain-0.27.0 → datachain-0.28.1}/src/datachain/studio.py

@@ -8,6 +8,7 @@ import dateparser
 import tabulate
 
 from datachain.config import Config, ConfigLevel
+from datachain.data_storage.job import JobStatus
 from datachain.dataset import QUERY_DATASET_PREFIX, parse_dataset_name
 from datachain.error import DataChainError
 from datachain.remote.studio import StudioClient
@@ -20,6 +21,8 @@ POST_LOGIN_MESSAGE = (
     "Once you've logged in, return here "
     "and you'll be ready to start using DataChain with Studio."
 )
+RETRY_MAX_TIMES = 10
+RETRY_SLEEP_SEC = 1
 
 
 def process_jobs_args(args: "Namespace"):
@@ -46,6 +49,7 @@ def process_jobs_args(args: "Namespace"):
             args.cluster,
             args.start_time,
             args.cron,
+            args.no_wait,
         )
 
     if args.cmd == "cancel":
@@ -287,17 +291,34 @@ def parse_start_time(start_time_str: Optional[str]) -> Optional[str]:
 def show_logs_from_client(client, job_id):
     # Sync usage
     async def _run():
+        retry_count = 0
         latest_status = None
-        async for message in client.tail_job_logs(job_id):
-            if "logs" in message:
-                for log in message["logs"]:
-                    print(log["message"], end="")
-            elif "job" in message:
-                latest_status = message["job"]["status"]
-                print(f"\n>>>> Job is now in {latest_status} status.")
+        processed_statuses = set()
+        while True:
+            async for message in client.tail_job_logs(job_id):
+                if "logs" in message:
+                    for log in message["logs"]:
+                        print(log["message"], end="")
+                elif "job" in message:
+                    latest_status = message["job"]["status"]
+                    if latest_status in processed_statuses:
+                        continue
+                    processed_statuses.add(latest_status)
+                    print(f"\n>>>> Job is now in {latest_status} status.")
+
+            try:
+                if retry_count > RETRY_MAX_TIMES or (
+                    latest_status and JobStatus[latest_status].finished()
+                ):
+                    break
+                await asyncio.sleep(RETRY_SLEEP_SEC)
+                retry_count += 1
+            except KeyError:
+                pass
+
         return latest_status
 
-    latest_status = asyncio.run(_run())
+    final_status = asyncio.run(_run())
 
     response = client.dataset_job_versions(job_id)
     if not response.ok:
@@ -314,9 +335,9 @@ def show_logs_from_client(client, job_id):
 
     exit_code_by_status = {
         "FAILED": 1,
-        "CANCELLED": 2,
+        "CANCELED": 2,
     }
-    return exit_code_by_status.get(latest_status.upper(), 0) if latest_status else 0
+    return exit_code_by_status.get(final_status.upper(), 0) if final_status else 0
 
 
 def create_job(
@@ -334,6 +355,7 @@ def create_job(
     cluster: Optional[str] = None,
     start_time: Optional[str] = None,
     cron: Optional[str] = None,
+    no_wait: Optional[bool] = False,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -388,7 +410,7 @@ def create_job(
     print("Open the job in Studio at", response.data.get("job", {}).get("url"))
     print("=" * 40)
 
-    return show_logs_from_client(client, job_id)
+    return 0 if no_wait else show_logs_from_client(client, job_id)
 
 
 def upload_files(client: StudioClient, files: list[str]) -> list[str]:

{datachain-0.27.0 → datachain-0.28.1}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.27.0
+Version: 0.28.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -45,7 +45,7 @@ Requires-Dist: datamodel-code-generator>=0.25
 Requires-Dist: Pillow<12,>=10.0.0
 Requires-Dist: msgpack<2,>=1.0.4
 Requires-Dist: psutil
-Requires-Dist: huggingface_hub
+Requires-Dist: huggingface_hub<0.34.0
 Requires-Dist: iterative-telemetry>=0.0.10
 Requires-Dist: platformdirs
 Requires-Dist: dvc-studio-client<1,>=0.21

{datachain-0.27.0 → datachain-0.28.1}/src/datachain.egg-info/requires.txt

@@ -26,7 +26,7 @@ datamodel-code-generator>=0.25
 Pillow<12,>=10.0.0
 msgpack<2,>=1.0.4
 psutil
-huggingface_hub
+huggingface_hub<0.34.0
 iterative-telemetry>=0.0.10
 platformdirs
 dvc-studio-client<1,>=0.21

{datachain-0.27.0 → datachain-0.28.1}/tests/func/test_datachain.py

@@ -9,7 +9,7 @@ import uuid
 from collections.abc import Iterator
 from datetime import datetime, timedelta, timezone
 from pathlib import Path, PurePosixPath
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import numpy as np
 import pandas as pd
@@ -358,15 +358,24 @@ def test_to_storage(
     file_type,
     num_threads,
 ):
+    mapper = Mock(side_effect=lambda file_path: len(file_path))
+
     ctc = cloud_test_catalog
     df = dc.read_storage(ctc.src_uri, type=file_type, session=test_session)
     if use_map:
-        df.settings(cache=use_cache).map(
-            res=lambda file: file.export(tmp_dir / "output", placement=placement)
-        ).exec()
+        (
+            df.settings(cache=use_cache)
+            .map(mapper, params=["file.path"], output={"path_len": int})
+            .map(res=lambda file: file.export(tmp_dir / "output", placement=placement))
+            .exec()
+        )
     else:
-        df.settings(cache=use_cache).to_storage(
-            tmp_dir / "output", placement=placement, num_threads=num_threads
+        (
+            df.settings(cache=use_cache)
+            .map(mapper, params=["file.path"], output={"path_len": int})
+            .to_storage(
+                tmp_dir / "output", placement=placement, num_threads=num_threads
+            )
         )
 
     expected = {
@@ -387,6 +396,8 @@ def test_to_storage(
         with open(tmp_dir / "output" / file_path) as f:
             assert f.read() == expected[file.name]
 
+    assert mapper.call_count == len(expected)
+
 
 @pytest.mark.parametrize("use_cache", [True, False])
 def test_export_images_files(test_session, tmp_dir, tmp_path, use_cache):

{datachain-0.27.0 → datachain-0.28.1}/tests/test_cli_studio.py

@@ -491,7 +491,7 @@ def test_studio_run_task(capsys, mocker, tmp_dir, studio_token):
 
 
 @pytest.mark.parametrize(
-    "status,expected_exit_code", [("FAILED", 1), ("CANCELLED", 2), ("COMPLETED", 0)]
+    "status,expected_exit_code", [("FAILED", 1), ("CANCELED", 2), ("COMPLETE", 0)]
 )
 def test_studio_run_non_zero_exit_code(
     capsys, mocker, tmp_dir, status, expected_exit_code

{datachain-0.27.0 → datachain-0.28.1}/tests/unit/lib/test_file.py

@@ -7,7 +7,7 @@ from fsspec.implementations.local import LocalFileSystem
 from PIL import Image
 
 from datachain.catalog import Catalog
-from datachain.lib.file import File, FileError, ImageFile, TextFile, resolve
+from datachain.lib.file import Audio, File, FileError, ImageFile, TextFile, resolve
 
 
 def create_file(source: str):
@@ -409,3 +409,49 @@ def test_path_normalized(path, expected, raises):
             file.get_path_normalized()
     else:
         assert file.get_path_normalized() == expected
+
+
+def test_file_rebase_method():
+    """Test File.rebase() method"""
+    file = File(source="s3://bucket", path="data/audio/file.wav")
+
+    # Basic rebase
+    result = file.rebase("s3://bucket/data/audio", "s3://output-bucket/waveforms")
+    assert result == "s3://output-bucket/waveforms/file.wav"
+
+    # With suffix and extension
+    result = file.rebase(
+        "s3://bucket/data/audio",
+        "s3://output-bucket/processed",
+        suffix="_ch1",
+        extension="npy",
+    )
+    assert result == "s3://output-bucket/processed/file_ch1.npy"
+
+
+def test_file_rebase_local_path():
+    """Test File.rebase() with local file paths"""
+    file = File(source="file://", path="/data/audio/folder/file.mp3")
+
+    result = file.rebase("file:///data/audio", "/output/processed")
+    assert result == "/output/processed/folder/file.mp3"
+
+
+def test_audio_get_channel_name():
+    # Test known channel configurations
+    assert Audio.get_channel_name(1, 0) == "Mono"
+    assert Audio.get_channel_name(2, 0) == "Left"
+    assert Audio.get_channel_name(2, 1) == "Right"
+    assert Audio.get_channel_name(4, 2) == "Y"  # Ambisonics
+    assert Audio.get_channel_name(6, 3) == "LFE"  # 5.1 surround
+    assert Audio.get_channel_name(8, 7) == "SR"  # 7.1 surround
+
+    # Test fallback for unknown configurations
+    assert Audio.get_channel_name(-1, 0) == "Ch1"
+    assert Audio.get_channel_name(3, 0) == "Ch1"
+    assert Audio.get_channel_name(5, 4) == "Ch5"
+    assert Audio.get_channel_name(10, 9) == "Ch10"
+
+    # Test out of range indices
+    assert Audio.get_channel_name(2, 5) == "Ch6"
+    assert Audio.get_channel_name(1, 1) == "Ch2"