datachain 0.30.5__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,13 +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 pyarrow as pa
-import ujson as json
 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
@@ -44,10 +44,10 @@ class ArrowGenerator(Generator):
 
     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,
     ):
         """
@@ -112,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:
@@ -141,7 +141,7 @@ class ArrowGenerator(Generator):
     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:
@@ -149,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
@@ -180,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.
@@ -205,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())
@@ -243,7 +246,7 @@ 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(f"ArrowDataModel_{column}", type_dict)
     if pa.types.is_map(col_type):
@@ -257,7 +260,7 @@ def arrow_type_mapper(col_type: pa.DataType, column: str = "") -> type:  # noqa:
 
 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
 
@@ -266,7 +269,7 @@ def _get_hf_schema(
     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 = json.loads(

datachain/lib/audio.py

@@ -1,5 +1,6 @@
 import posixpath
-from typing import TYPE_CHECKING, Optional, Union
+import re
+from typing import TYPE_CHECKING
 
 from datachain.lib.file import FileError
 
@@ -9,7 +10,7 @@ if TYPE_CHECKING:
     from datachain.lib.file import Audio, AudioFile, File
 
 try:
-    import torchaudio
+    import soundfile as sf
 except ImportError as exc:
     raise ImportError(
         "Missing dependencies for processing audio.\n"
@@ -18,7 +19,7 @@ except ImportError as exc:
     ) from exc
 
 
-def audio_info(file: "Union[File, AudioFile]") -> "Audio":
+def audio_info(file: "File | AudioFile") -> "Audio":
     """Extract metadata like sample rate, channels, duration, and format."""
     from datachain.lib.file import Audio
 
@@ -26,18 +27,25 @@ def audio_info(file: "Union[File, AudioFile]") -> "Audio":
 
     try:
         with file.open() as f:
-            info = torchaudio.info(f)
+            info = sf.info(f)
+
+            sample_rate = int(info.samplerate)
+            channels = int(info.channels)
+            frames = int(info.frames)
+            duration = float(info.duration)
 
-            sample_rate = int(info.sample_rate)
-            channels = int(info.num_channels)
-            frames = int(info.num_frames)
-            duration = float(frames / sample_rate) if sample_rate > 0 else 0.0
+            # soundfile provides format and subtype
+            if info.format:
+                format_name = info.format.lower()
+            else:
+                format_name = file.get_file_ext().lower()
 
-            codec_name = getattr(info, "encoding", "")
-            file_ext = file.get_file_ext().lower()
-            format_name = _encoding_to_format(codec_name, file_ext)
+            if not format_name:
+                format_name = "unknown"
+            codec_name = info.subtype if info.subtype else ""
 
-            bits_per_sample = getattr(info, "bits_per_sample", 0)
+            # 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
             )
@@ -58,48 +66,43 @@ def audio_info(file: "Union[File, AudioFile]") -> "Audio":
     )
 
 
-def _encoding_to_format(encoding: str, file_ext: str) -> str:
+def _get_bits_per_sample(subtype: str) -> int:
     """
-    Map torchaudio encoding to a format name.
+    Map soundfile subtype to bits per sample.
 
     Args:
-        encoding: The encoding string from torchaudio.info()
-        file_ext: The file extension as a fallback
+        subtype: The subtype string from soundfile
 
     Returns:
-        Format name as a string
+        Bits per sample, or 0 if unknown
     """
-    # Direct mapping for formats that match exactly
-    encoding_map = {
-        "FLAC": "flac",
-        "MP3": "mp3",
-        "VORBIS": "ogg",
-        "AMR_WB": "amr",
-        "AMR_NB": "amr",
-        "OPUS": "opus",
-        "GSM": "gsm",
+    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 encoding in encoding_map:
-        return encoding_map[encoding]
+    if subtype in pcm_bits:
+        return pcm_bits[subtype]
 
-    # For PCM variants, use file extension to determine format
-    if encoding.startswith("PCM_"):
-        # Common PCM formats by extension
-        pcm_formats = {
-            "wav": "wav",
-            "aiff": "aiff",
-            "au": "au",
-            "raw": "raw",
-        }
-        return pcm_formats.get(file_ext, "wav")  # Default to wav for PCM
+    # 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))
 
-    # Fallback to file extension if encoding is unknown
-    return file_ext if file_ext else "unknown"
+    return 0
 
 
 def audio_to_np(
-    audio: "AudioFile", start: float = 0, duration: Optional[float] = None
+    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)."""
@@ -114,27 +117,27 @@ def audio_to_np(
 
     try:
         with audio.open() as f:
-            info = torchaudio.info(f)
-            sample_rate = info.sample_rate
+            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
-            # This is important to ensure we read from the correct position later
             f.seek(0)
 
-            waveform, sr = torchaudio.load(
-                f, frame_offset=frame_offset, num_frames=num_frames
+            # 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",
             )
 
-            audio_np = waveform.numpy()
-
-            if audio_np.shape[0] > 1:
-                audio_np = audio_np.T
-            else:
-                audio_np = audio_np.squeeze()
-
+            # 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(
@@ -146,17 +149,15 @@ def audio_to_bytes(
     audio: "AudioFile",
     format: str = "wav",
     start: float = 0,
-    duration: Optional[float] = None,
+    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."""
-    y, sr = audio_to_np(audio, start, duration)
-
     import io
 
-    import soundfile as sf
+    y, sr = audio_to_np(audio, start, duration)
 
     buffer = io.BytesIO()
     sf.write(buffer, y, sr, format=format)
@@ -166,9 +167,9 @@ def audio_to_bytes(
 def save_audio(
     audio: "AudioFile",
     output: str,
-    format: Optional[str] = None,
+    format: str | None = None,
     start: float = 0,
-    end: Optional[float] = None,
+    end: float | None = None,
 ) -> "AudioFile":
     """Save audio file or extract fragment to specified format.
 

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/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)