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

datachain/json.py

@@ -0,0 +1,138 @@
+"""DataChain JSON utilities.
+
+This module wraps :mod:`ujson` so we can guarantee consistent handling
+of values that the encoder does not support out of the box (for example
+``datetime`` objects or ``bytes``).
+All code inside DataChain should import this module instead of using
+:mod:`ujson` directly.
+"""
+
+import datetime as _dt
+import json as _json
+import uuid as _uuid
+from collections.abc import Callable
+from typing import Any
+
+import ujson as _ujson
+
+__all__ = [
+    "JSONDecodeError",
+    "dump",
+    "dumps",
+    "load",
+    "loads",
+]
+
+JSONDecodeError = (_ujson.JSONDecodeError, _json.JSONDecodeError)
+
+_SENTINEL = object()
+_Default = Callable[[Any], Any]
+DEFAULT_PREVIEW_BYTES = 1024
+
+
+# To make it looks like Pydantic's ISO format with 'Z' for UTC
+# It is minor but nice to have consistency
+def _format_datetime(value: _dt.datetime) -> str:
+    iso = value.isoformat()
+
+    offset = value.utcoffset()
+    if value.tzinfo is None or offset is None:
+        return iso
+
+    if offset == _dt.timedelta(0) and iso.endswith(("+00:00", "-00:00")):
+        return iso[:-6] + "Z"
+
+    return iso
+
+
+def _format_time(value: _dt.time) -> str:
+    iso = value.isoformat()
+
+    offset = value.utcoffset()
+    if value.tzinfo is None or offset is None:
+        return iso
+
+    if offset == _dt.timedelta(0) and iso.endswith(("+00:00", "-00:00")):
+        return iso[:-6] + "Z"
+
+    return iso
+
+
+def _coerce(value: Any, serialize_bytes: bool) -> Any:
+    """Return a JSON-serializable representation for supported extra types."""
+
+    if isinstance(value, _dt.datetime):
+        return _format_datetime(value)
+    if isinstance(value, _dt.date):
+        return value.isoformat()
+    if isinstance(value, _dt.time):
+        return _format_time(value)
+    if isinstance(value, _uuid.UUID):
+        return str(value)
+    if serialize_bytes and isinstance(value, (bytes, bytearray)):
+        return list(bytes(value)[:DEFAULT_PREVIEW_BYTES])
+    return _SENTINEL
+
+
+def _base_default(value: Any, serialize_bytes: bool) -> Any:
+    converted = _coerce(value, serialize_bytes)
+    if converted is not _SENTINEL:
+        return converted
+    raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable")
+
+
+def _build_default(user_default: _Default | None, serialize_bytes: bool) -> _Default:
+    if user_default is None:
+        return lambda value: _base_default(value, serialize_bytes)
+
+    def combined(value: Any) -> Any:
+        converted = _coerce(value, serialize_bytes)
+        if converted is not _SENTINEL:
+            return converted
+        return user_default(value)
+
+    return combined
+
+
+def dumps(
+    obj: Any,
+    *,
+    default: _Default | None = None,
+    serialize_bytes: bool = False,
+    **kwargs: Any,
+) -> str:
+    """Serialize *obj* to a JSON-formatted ``str``."""
+
+    if serialize_bytes:
+        return _json.dumps(obj, default=_build_default(default, True), **kwargs)
+
+    return _ujson.dumps(obj, default=_build_default(default, False), **kwargs)
+
+
+def dump(
+    obj: Any,
+    fp,
+    *,
+    default: _Default | None = None,
+    serialize_bytes: bool = False,
+    **kwargs: Any,
+) -> None:
+    """Serialize *obj* as a JSON formatted stream to *fp*."""
+
+    if serialize_bytes:
+        _json.dump(obj, fp, default=_build_default(default, True), **kwargs)
+        return
+
+    _ujson.dump(obj, fp, default=_build_default(default, False), **kwargs)
+
+
+def loads(s: str | bytes | bytearray, **kwargs: Any) -> Any:
+    """Deserialize *s* to a Python object."""
+
+    return _ujson.loads(s, **kwargs)
+
+
+def load(fp, **kwargs: Any) -> Any:
+    """Deserialize JSON content from *fp* to a Python object."""
+
+    return loads(fp.read(), **kwargs)

datachain/lib/arrow.py

@@ -1,12 +1,13 @@
 from collections.abc import Sequence
 from itertools import islice
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
-import orjson
 import pyarrow as pa
+from pyarrow._csv import ParseOptions
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm.auto import tqdm
 
+from datachain import json
 from datachain.fs.reference import ReferenceFileSystem
 from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import ArrowRow, File
@@ -26,15 +27,27 @@ if TYPE_CHECKING:
 DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY = b"DataChain SignalSchema"
 
 
+def fix_pyarrow_format(format, parse_options=None):
+    # Re-init invalid row handler: https://issues.apache.org/jira/browse/ARROW-17641
+    if (
+        format
+        and isinstance(format, CsvFileFormat)
+        and parse_options
+        and isinstance(parse_options, ParseOptions)
+    ):
+        format.parse_options = parse_options
+    return format
+
+
 class ArrowGenerator(Generator):
     DEFAULT_BATCH_SIZE = 2**17  # same as `pyarrow._dataset._DEFAULT_BATCH_SIZE`
 
     def __init__(
         self,
-        input_schema: Optional["pa.Schema"] = None,
-        output_schema: Optional[type["BaseModel"]] = None,
+        input_schema: pa.Schema | None = None,
+        output_schema: type["BaseModel"] | None = None,
         source: bool = True,
-        nrows: Optional[int] = None,
+        nrows: int | None = None,
         **kwargs,
     ):
         """
@@ -53,6 +66,7 @@ class ArrowGenerator(Generator):
         self.output_schema = output_schema
         self.source = source
         self.nrows = nrows
+        self.parse_options = kwargs.pop("parse_options", None)
         self.kwargs = kwargs
 
     def process(self, file: File):
@@ -62,9 +76,13 @@ class ArrowGenerator(Generator):
             fs_path = file.path
             fs = ReferenceFileSystem({fs_path: [cache_path]})
         else:
-            fs, fs_path = file.get_fs(), file.get_path()
+            fs, fs_path = file.get_fs(), file.get_fs_path()
+
+        kwargs = self.kwargs
+        if format := kwargs.get("format"):
+            kwargs["format"] = fix_pyarrow_format(format, self.parse_options)
 
-        ds = dataset(fs_path, schema=self.input_schema, filesystem=fs, **self.kwargs)
+        ds = dataset(fs_path, schema=self.input_schema, filesystem=fs, **kwargs)
 
         hf_schema = _get_hf_schema(ds.schema)
         use_datachain_schema = (
@@ -94,7 +112,7 @@ class ArrowGenerator(Generator):
         record: dict[str, Any],
         file: File,
         index: int,
-        hf_schema: Optional[tuple["Features", dict[str, "DataType"]]],
+        hf_schema: tuple["Features", dict[str, "DataType"]] | None,
         use_datachain_schema: bool,
     ):
         if use_datachain_schema and self.output_schema:
@@ -108,13 +126,22 @@ class ArrowGenerator(Generator):
             if isinstance(kwargs.get("format"), CsvFileFormat):
                 kwargs["format"] = "csv"
             arrow_file = ArrowRow(file=file, index=index, kwargs=kwargs)
+
+            if self.output_schema and hasattr(vals[0], "source"):
+                # if we are reading parquet file written by datachain it might have
+                # source inside of it already, so we should not duplicate it, instead
+                # we are re-creating it of the self.source flag
+                vals[0].source = arrow_file  # type: ignore[attr-defined]
+
+                return vals
             return [arrow_file, *vals]
+
         return vals
 
     def _process_non_datachain_record(
         self,
         record: dict[str, Any],
-        hf_schema: Optional[tuple["Features", dict[str, "DataType"]]],
+        hf_schema: tuple["Features", dict[str, "DataType"]] | None,
     ):
         vals = list(record.values())
         if not self.output_schema:
@@ -122,7 +149,9 @@ class ArrowGenerator(Generator):
 
         fields = self.output_schema.model_fields
         vals_dict = {}
-        for i, ((field, field_info), val) in enumerate(zip(fields.items(), vals)):
+        for i, ((field, field_info), val) in enumerate(
+            zip(fields.items(), vals, strict=False)
+        ):
             anno = field_info.annotation
             if hf_schema:
                 from datachain.lib.hf import convert_feature
@@ -137,9 +166,13 @@ class ArrowGenerator(Generator):
 
 
 def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
+    parse_options = kwargs.pop("parse_options", None)
+    if format := kwargs.get("format"):
+        kwargs["format"] = fix_pyarrow_format(format, parse_options)
+
     schemas = []
-    for file in chain.collect("file"):
-        ds = dataset(file.get_path(), filesystem=file.get_fs(), **kwargs)  # type: ignore[union-attr]
+    for (file,) in chain.to_iter("file"):
+        ds = dataset(file.get_fs_path(), filesystem=file.get_fs(), **kwargs)  # type: ignore[union-attr]
         schemas.append(ds.schema)
     if not schemas:
         raise ValueError(
@@ -149,7 +182,7 @@ def infer_schema(chain: "DataChain", **kwargs) -> pa.Schema:
 
 
 def schema_to_output(
-    schema: pa.Schema, col_names: Optional[Sequence[str]] = None
+    schema: pa.Schema, col_names: Sequence[str] | None = None
 ) -> tuple[dict[str, type], list[str]]:
     """
     Generate UDF output schema from pyarrow schema.
@@ -174,14 +207,15 @@ def schema_to_output(
     hf_schema = _get_hf_schema(schema)
     if hf_schema:
         return {
-            column: hf_type for hf_type, column in zip(hf_schema[1].values(), col_names)
+            column: hf_type
+            for hf_type, column in zip(hf_schema[1].values(), col_names, strict=False)
         }, list(normalized_col_dict.values())
 
     output = {}
-    for field, column in zip(schema, col_names):
+    for field, column in zip(schema, col_names, strict=False):
         dtype = arrow_type_mapper(field.type, column)
         if field.nullable and not ModelStore.is_pydantic(dtype):
-            dtype = Optional[dtype]  # type: ignore[assignment]
+            dtype = dtype | None  # type: ignore[assignment]
         output[column] = dtype
 
     return output, list(normalized_col_dict.values())
@@ -212,31 +246,33 @@ def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa:
         for field in col_type:
             dtype = arrow_type_mapper(field.type, field.name)
             if field.nullable and not ModelStore.is_pydantic(dtype):
-                dtype = Optional[dtype]  # type: ignore[assignment]
+                dtype = dtype | None  # type: ignore[assignment]
             type_dict[field.name] = dtype
-        return dict_to_data_model(column, type_dict)
+        return dict_to_data_model(f"ArrowDataModel_{column}", type_dict)
     if pa.types.is_map(col_type):
         return dict
     if isinstance(col_type, pa.lib.DictionaryType):
         return arrow_type_mapper(col_type.value_type)  # type: ignore[return-value]
+    if pa.types.is_null(col_type):
+        return str  # use strings for null columns
     raise TypeError(f"{col_type!r} datatypes not supported, column: {column}")
 
 
 def _get_hf_schema(
     schema: "pa.Schema",
-) -> Optional[tuple["Features", dict[str, "DataType"]]]:
+) -> tuple["Features", dict[str, "DataType"]] | None:
     if schema.metadata and b"huggingface" in schema.metadata:
         from datachain.lib.hf import get_output_schema, schema_from_arrow
 
         features = schema_from_arrow(schema)
-        return features, get_output_schema(features)
+        return features, get_output_schema(features)[0]
     return None
 
 
-def _get_datachain_schema(schema: "pa.Schema") -> Optional[SignalSchema]:
+def _get_datachain_schema(schema: "pa.Schema") -> SignalSchema | None:
     """Return a restored SignalSchema from parquet metadata, if any is found."""
     if schema.metadata and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in schema.metadata:
-        serialized_signal_schema = orjson.loads(
+        serialized_signal_schema = json.loads(
             schema.metadata[DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY]
         )
         return SignalSchema.deserialize(serialized_signal_schema)

datachain/lib/audio.py

@@ -0,0 +1,245 @@
+import posixpath
+import re
+from typing import TYPE_CHECKING
+
+from datachain.lib.file import FileError
+
+if TYPE_CHECKING:
+    from numpy import ndarray
+
+    from datachain.lib.file import Audio, AudioFile, File
+
+try:
+    import soundfile as sf
+except ImportError as exc:
+    raise ImportError(
+        "Missing dependencies for processing audio.\n"
+        "To install run:\n\n"
+        "  pip install 'datachain[audio]'\n"
+    ) from exc
+
+
+def audio_info(file: "File | AudioFile") -> "Audio":
+    """Extract metadata like sample rate, channels, duration, and format."""
+    from datachain.lib.file import Audio
+
+    file = file.as_audio_file()
+
+    try:
+        with file.open() as f:
+            info = sf.info(f)
+
+            sample_rate = int(info.samplerate)
+            channels = int(info.channels)
+            frames = int(info.frames)
+            duration = float(info.duration)
+
+            # soundfile provides format and subtype
+            if info.format:
+                format_name = info.format.lower()
+            else:
+                format_name = file.get_file_ext().lower()
+
+            if not format_name:
+                format_name = "unknown"
+            codec_name = info.subtype if info.subtype else ""
+
+            # Calculate bit rate from subtype
+            bits_per_sample = _get_bits_per_sample(info.subtype)
+            bit_rate = (
+                bits_per_sample * sample_rate * channels if bits_per_sample > 0 else -1
+            )
+
+    except Exception as exc:
+        raise FileError(
+            "unable to extract metadata from audio file", file.source, file.path
+        ) from exc
+
+    return Audio(
+        sample_rate=sample_rate,
+        channels=channels,
+        duration=duration,
+        samples=frames,
+        format=format_name,
+        codec=codec_name,
+        bit_rate=bit_rate,
+    )
+
+
+def _get_bits_per_sample(subtype: str) -> int:
+    """
+    Map soundfile subtype to bits per sample.
+
+    Args:
+        subtype: The subtype string from soundfile
+
+    Returns:
+        Bits per sample, or 0 if unknown
+    """
+    if not subtype:
+        return 0
+
+    # Common PCM and floating-point subtypes
+    pcm_bits = {
+        "PCM_16": 16,
+        "PCM_24": 24,
+        "PCM_32": 32,
+        "PCM_S8": 8,
+        "PCM_U8": 8,
+        "FLOAT": 32,
+        "DOUBLE": 64,
+    }
+
+    if subtype in pcm_bits:
+        return pcm_bits[subtype]
+
+    # Handle variants such as PCM_S16LE, PCM_F32LE, etc.
+    match = re.search(r"PCM_(?:[A-Z]*?)(\d+)", subtype)
+    if match:
+        return int(match.group(1))
+
+    return 0
+
+
+def audio_to_np(
+    audio: "AudioFile", start: float = 0, duration: float | None = None
+) -> "tuple[ndarray, int]":
+    """Load audio fragment as numpy array.
+    Multi-channel audio is transposed to (samples, channels)."""
+    if start < 0:
+        raise ValueError("start must be a non-negative float")
+
+    if duration is not None and duration <= 0:
+        raise ValueError("duration must be a positive float")
+
+    if hasattr(audio, "as_audio_file"):
+        audio = audio.as_audio_file()
+
+    try:
+        with audio.open() as f:
+            info = sf.info(f)
+            sample_rate = info.samplerate
+
+            frame_offset = int(start * sample_rate)
+            num_frames = int(duration * sample_rate) if duration is not None else -1
+
+            # Reset file pointer to the beginning
+            f.seek(0)
+
+            # Read audio data with offset and frame count
+            audio_np, sr = sf.read(
+                f,
+                start=frame_offset,
+                frames=num_frames,
+                always_2d=False,
+                dtype="float32",
+            )
+
+            # soundfile returns shape (frames,) for mono or
+            # (frames, channels) for multi-channel
+            # We keep this format as it matches expected output
+            return audio_np, int(sr)
+    except Exception as exc:
+        raise FileError(
+            "unable to read audio fragment", audio.source, audio.path
+        ) from exc
+
+
+def audio_to_bytes(
+    audio: "AudioFile",
+    format: str = "wav",
+    start: float = 0,
+    duration: float | None = None,
+) -> bytes:
+    """Convert audio to bytes using soundfile.
+
+    If duration is None, converts from start to end of file.
+    If start is 0 and duration is None, converts entire file."""
+    import io
+
+    y, sr = audio_to_np(audio, start, duration)
+
+    buffer = io.BytesIO()
+    sf.write(buffer, y, sr, format=format)
+    return buffer.getvalue()
+
+
+def save_audio(
+    audio: "AudioFile",
+    output: str,
+    format: str | None = None,
+    start: float = 0,
+    end: float | None = None,
+) -> "AudioFile":
+    """Save audio file or extract fragment to specified format.
+
+    Args:
+        audio: Source AudioFile object
+        output: Output directory path
+        format: Output format ('wav', 'mp3', etc). Defaults to source format
+        start: Start time in seconds (>= 0). Defaults to 0
+        end: End time in seconds. If None, extracts to end of file
+
+    Returns:
+        AudioFile: New audio file with format conversion/extraction applied
+
+    Examples:
+        save_audio(audio, "/path", "mp3")                       # Entire file to MP3
+        save_audio(audio, "s3://bucket/path", "wav", start=2.5) # From 2.5s to end
+        save_audio(audio, "/path", "flac", start=1, end=3)      # Extract 1-3s fragment
+    """
+    if format is None:
+        format = audio.get_file_ext()
+
+    # Validate start time
+    if start < 0:
+        raise ValueError(
+            f"Can't save audio for '{audio.path}', "
+            f"start time must be non-negative: {start:.3f}"
+        )
+
+    # Handle full file conversion when end is None and start is 0
+    if end is None and start == 0:
+        output_file = posixpath.join(output, f"{audio.get_file_stem()}.{format}")
+        try:
+            audio_bytes = audio_to_bytes(audio, format, start=0, duration=None)
+        except Exception as exc:
+            raise FileError(
+                "unable to convert audio file", audio.source, audio.path
+            ) from exc
+    elif end is None:
+        # Extract from start to end of file
+        output_file = posixpath.join(
+            output, f"{audio.get_file_stem()}_{int(start * 1000):06d}_end.{format}"
+        )
+        try:
+            audio_bytes = audio_to_bytes(audio, format, start=start, duration=None)
+        except Exception as exc:
+            raise FileError(
+                "unable to save audio fragment", audio.source, audio.path
+            ) from exc
+    else:
+        # Fragment extraction mode with specific end time
+        if end < 0 or start >= end:
+            raise ValueError(
+                f"Can't save audio for '{audio.path}', "
+                f"invalid time range: ({start:.3f}, {end:.3f})"
+            )
+
+        duration = end - start
+        start_ms = int(start * 1000)
+        end_ms = int(end * 1000)
+        output_file = posixpath.join(
+            output, f"{audio.get_file_stem()}_{start_ms:06d}_{end_ms:06d}.{format}"
+        )
+
+        try:
+            audio_bytes = audio_to_bytes(audio, format, start, duration)
+        except Exception as exc:
+            raise FileError(
+                "unable to save audio fragment", audio.source, audio.path
+            ) from exc
+
+    from datachain.lib.file import AudioFile
+
+    return AudioFile.upload(audio_bytes, output_file, catalog=audio._catalog)

datachain/lib/clip.py

@@ -1,5 +1,6 @@
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, Union
 
 import torch
 from transformers.modeling_utils import PreTrainedModel
@@ -32,28 +33,28 @@ def _get_encoder(model: Any, type: Literal["image", "text"]) -> Callable:
 
 
 def clip_similarity_scores(
-    images: Union[None, "Image.Image", list["Image.Image"]],
-    text: Union[None, str, list[str]],
+    images: Union["Image.Image", list["Image.Image"]] | None,
+    text: str | list[str] | None,
     model: Any,
     preprocess: Callable,
     tokenizer: Callable,
     prob: bool = False,
     image_to_text: bool = True,
-    device: Optional[Union[str, torch.device]] = None,
+    device: str | torch.device | None = None,
 ) -> list[list[float]]:
     """
     Calculate CLIP similarity scores between one or more images and/or text.
 
     Parameters:
-        images : Images to use as inputs.
-        text : Text to use as inputs.
-        model : Model from clip or open_clip packages.
-        preprocess : Image preprocessor to apply.
-        tokenizer : Text tokenizer.
-        prob : Compute softmax probabilities.
-        image_to_text : Whether to compute for image-to-text or text-to-image. Ignored
-            if only one of images or text provided.
-        device : Device to use. Defaults is None - use model's device.
+        images: Images to use as inputs.
+        text: Text to use as inputs.
+        model: Model from clip or open_clip packages.
+        preprocess: Image preprocessor to apply.
+        tokenizer: Text tokenizer.
+        prob: Compute softmax probabilities.
+        image_to_text: Whether to compute for image-to-text or text-to-image. Ignored
+            if only one of the images or text provided.
+        device: Device to use. Default is None - use model's device.
 
 
     Example:

datachain/lib/convert/flatten.py

@@ -6,12 +6,14 @@ from datachain.lib.model_store import ModelStore
 
 
 def flatten(obj: BaseModel) -> tuple:
-    return tuple(_flatten_fields_values(obj.model_fields, obj))
+    return tuple(_flatten_fields_values(type(obj).model_fields, obj))
 
 
 def flatten_list(obj_list: list[BaseModel]) -> tuple:
     return tuple(
-        val for obj in obj_list for val in _flatten_fields_values(obj.model_fields, obj)
+        val
+        for obj in obj_list
+        for val in _flatten_fields_values(type(obj).model_fields, obj)
     )
 
 
@@ -43,4 +45,4 @@ def _flatten_fields_values(fields: dict, obj: BaseModel) -> Generator:
 
 
 def _flatten(obj: BaseModel) -> tuple:
-    return tuple(_flatten_fields_values(obj.model_fields, obj))
+    return tuple(_flatten_fields_values(type(obj).model_fields, obj))

datachain/lib/convert/python_to_sql.py

@@ -1,6 +1,7 @@
 import inspect
 from datetime import datetime
 from enum import Enum
+from types import UnionType
 from typing import Annotated, Literal, Union, get_args, get_origin
 
 from pydantic import BaseModel
@@ -69,11 +70,12 @@ def python_to_sql(typ):  # noqa: PLR0911
     if inspect.isclass(orig) and issubclass(dict, orig):
         return JSON
 
-    if orig == Union:
+    if orig in (Union, UnionType):
         if len(args) == 2 and (type(None) in args):
-            return python_to_sql(args[0])
+            non_none_arg = args[0] if args[0] is not type(None) else args[1]
+            return python_to_sql(non_none_arg)
 
-        if _is_union_str_literal(orig, args):
+        if all(arg is str or get_origin(arg) in (Literal, LiteralEx) for arg in args):
             return String
 
         if _is_json_inside_union(orig, args):
@@ -95,7 +97,7 @@ def list_of_args_to_type(args) -> SQLType:
 
 
 def _is_json_inside_union(orig, args) -> bool:
-    if orig == Union and len(args) >= 2:
+    if orig in (Union, UnionType) and len(args) >= 2:
         # List in JSON: Union[dict, list[dict]]
         args_no_nones = [arg for arg in args if arg != type(None)]  # noqa: E721
         if len(args_no_nones) == 2:
@@ -109,9 +111,3 @@ def _is_json_inside_union(orig, args) -> bool:
         if any(inspect.isclass(arg) and issubclass(arg, BaseModel) for arg in args):
             return True
     return False
-
-
-def _is_union_str_literal(orig, args) -> bool:
-    if orig != Union:
-        return False
-    return all(arg is str or get_origin(arg) in (Literal, LiteralEx) for arg in args)