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

datachain/sql/sqlite/base.py

@@ -2,20 +2,20 @@ import logging
 import re
 import sqlite3
 import warnings
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
+from contextlib import closing
 from datetime import MAXYEAR, MINYEAR, datetime, timezone
 from functools import cache
 from types import MappingProxyType
-from typing import Callable, Optional
 
 import sqlalchemy as sa
-import ujson as json
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import case
 from sqlalchemy.sql.functions import func
 
+from datachain import json
 from datachain.sql.functions import (
     aggregate,
     array,
@@ -112,7 +112,10 @@ def setup():
     compiles(numeric.int_hash_64, "sqlite")(compile_int_hash_64)
     compiles(numeric.bit_hamming_distance, "sqlite")(compile_bit_hamming_distance)
 
-    if load_usearch_extension(sqlite3.connect(":memory:")):
+    with closing(sqlite3.connect(":memory:")) as _usearch_conn:
+        usearch_available = load_usearch_extension(_usearch_conn)
+
+    if usearch_available:
         compiles(array.cosine_distance, "sqlite")(compile_cosine_distance_ext)
         compiles(array.euclidean_distance, "sqlite")(compile_euclidean_distance_ext)
     else:
@@ -132,7 +135,7 @@ def run_compiler_hook(name):
 
 
 def functions_exist(
-    names: Iterable[str], connection: Optional[sqlite3.Connection] = None
+    names: Iterable[str], connection: sqlite3.Connection | None = None
 ) -> bool:
     """
     Returns True if all function names are defined for the given connection.
@@ -146,23 +149,34 @@ def functions_exist(
                 f"Found value of type {type(n).__name__}: {n!r}"
             )
 
+    close_connection = False
     if connection is None:
         connection = sqlite3.connect(":memory:")
+        close_connection = True
 
-    if not names:
-        return True
-    column1 = sa.column("column1", sa.String)
-    func_name_query = column1.not_in(
-        sa.select(sa.column("name", sa.String)).select_from(func.pragma_function_list())
-    )
-    query = (
-        sa.select(func.count() == 0)
-        .select_from(sa.values(column1).data([(n,) for n in names]))
-        .where(func_name_query)
-    )
-    comp = query.compile(dialect=sqlite_dialect)
-    args = (comp.string, comp.params) if comp.params else (comp.string,)
-    return bool(connection.execute(*args).fetchone()[0])
+    try:
+        if not names:
+            return True
+        column1 = sa.column("column1", sa.String)
+        func_name_query = column1.not_in(
+            sa.select(sa.column("name", sa.String)).select_from(
+                func.pragma_function_list()
+            )
+        )
+        query = (
+            sa.select(func.count() == 0)
+            .select_from(sa.values(column1).data([(n,) for n in names]))
+            .where(func_name_query)
+        )
+        comp = query.compile(dialect=sqlite_dialect)
+        if comp.params:
+            result = connection.execute(comp.string, comp.params)
+        else:
+            result = connection.execute(comp.string)
+        return bool(result.fetchone()[0])
+    finally:
+        if close_connection:
+            connection.close()
 
 
 def create_user_defined_sql_functions(connection):
@@ -201,9 +215,7 @@ def sqlite_int_hash_64(x: int) -> int:
 def sqlite_bit_hamming_distance(a: int, b: int) -> int:
     """Calculate the Hamming distance between two integers."""
     diff = (a & MAX_INT64) ^ (b & MAX_INT64)
-    if hasattr(diff, "bit_count"):
-        return diff.bit_count()
-    return bin(diff).count("1")
+    return diff.bit_count()
 
 
 def sqlite_byte_hamming_distance(a: str, b: str) -> int:
@@ -215,7 +227,7 @@ def sqlite_byte_hamming_distance(a: str, b: str) -> int:
     elif len(b) < len(a):
         diff = len(a) - len(b)
         a = a[: len(b)]
-    return diff + sum(c1 != c2 for c1, c2 in zip(a, b))
+    return diff + sum(c1 != c2 for c1, c2 in zip(a, b, strict=False))
 
 
 def register_user_defined_sql_functions() -> None:
@@ -470,7 +482,7 @@ def py_json_array_get_element(val, idx):
         return None
 
 
-def py_json_array_slice(val, offset: int, length: Optional[int] = None):
+def py_json_array_slice(val, offset: int, length: int | None = None):
     arr = json.loads(val)
     try:
         return json.dumps(
@@ -605,7 +617,7 @@ def compile_collect(element, compiler, **kwargs):
 
 
 @cache
-def usearch_sqlite_path() -> Optional[str]:
+def usearch_sqlite_path() -> str | None:
     try:
         import usearch
     except ImportError:

datachain/sql/sqlite/types.py

@@ -1,8 +1,8 @@
 import sqlite3
 
-import ujson as json
 from sqlalchemy import types
 
+from datachain import json
 from datachain.sql.types import TypeConverter, TypeReadConverter
 
 try:

datachain/sql/types.py

@@ -12,14 +12,15 @@ for sqlite we can use `sqlite.register_converter`
 ( https://docs.python.org/3/library/sqlite3.html#sqlite3.register_converter )
 """
 
+import numbers
 from datetime import datetime
 from types import MappingProxyType
 from typing import Any, Union
 
 import sqlalchemy as sa
-import ujson as jsonlib
 from sqlalchemy import TypeDecorator, types
 
+from datachain import json as jsonlib
 from datachain.lib.data_model import StandardType
 
 _registry: dict[str, "TypeConverter"] = {}
@@ -336,10 +337,28 @@ class Array(SQLType):
 
     @classmethod
     def from_dict(cls, d: dict[str, Any]) -> Union[type["SQLType"], "SQLType"]:
-        sub_t = NAME_TYPES_MAPPING[d["item_type"]["type"]].from_dict(  # type: ignore [attr-defined]
-            d["item_type"]
-        )
-        return cls(sub_t)
+        try:
+            array_item = d["item_type"]
+        except KeyError as e:
+            raise ValueError("Array type must have 'item_type' field") from e
+
+        if not isinstance(array_item, dict):
+            raise TypeError("Array 'item_type' field must be a dictionary")
+
+        try:
+            item_type = array_item["type"]
+        except KeyError as e:
+            raise ValueError("Array 'item_type' must have 'type' field") from e
+
+        try:
+            sub_t = NAME_TYPES_MAPPING[item_type]
+        except KeyError as e:
+            raise ValueError(f"Array item type '{item_type}' is not supported") from e
+
+        try:
+            return cls(sub_t.from_dict(d["item_type"]))  # type: ignore [attr-defined]
+        except KeyError as e:
+            raise ValueError(f"Array item type '{item_type}' is not supported") from e
 
     @staticmethod
     def default_value(dialect):
@@ -427,6 +446,18 @@ class TypeReadConverter:
         return value
 
     def boolean(self, value):
+        if value is None or isinstance(value, bool):
+            return value
+
+        if isinstance(value, numbers.Integral):
+            return bool(value)
+        if isinstance(value, str):
+            normalized = value.strip().lower()
+            if normalized in {"true", "t", "yes", "y", "1"}:
+                return True
+            if normalized in {"false", "f", "no", "n", "0"}:
+                return False
+
         return value
 
     def int(self, value):

datachain/studio.py

@@ -1,8 +1,9 @@
 import asyncio
 import os
 import sys
+import warnings
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 import dateparser
 import tabulate
@@ -175,7 +176,7 @@ def token():
     print(token)
 
 
-def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
+def list_datasets(team: str | None = None, name: str | None = None):
     def ds_full_name(ds: dict) -> str:
         return (
             f"{ds['project']['namespace']['name']}.{ds['project']['name']}.{ds['name']}"
@@ -206,7 +207,7 @@ def list_datasets(team: Optional[str] = None, name: Optional[str] = None):
             yield (full_name, version)
 
 
-def list_dataset_versions(team: Optional[str] = None, name: str = ""):
+def list_dataset_versions(team: str | None = None, name: str = ""):
     client = StudioClient(team=team)
 
     namespace_name, project_name, name = parse_dataset_name(name)
@@ -226,13 +227,13 @@ def list_dataset_versions(team: Optional[str] = None, name: str = ""):
 
 
 def edit_studio_dataset(
-    team_name: Optional[str],
+    team_name: str | None,
     name: str,
     namespace: str,
     project: str,
-    new_name: Optional[str] = None,
-    description: Optional[str] = None,
-    attrs: Optional[list[str]] = None,
+    new_name: str | None = None,
+    description: str | None = None,
+    attrs: list[str] | None = None,
 ):
     client = StudioClient(team=team_name)
     response = client.edit_dataset(
@@ -245,12 +246,12 @@ def edit_studio_dataset(
 
 
 def remove_studio_dataset(
-    team_name: Optional[str],
+    team_name: str | None,
     name: str,
     namespace: str,
     project: str,
-    version: Optional[str] = None,
-    force: Optional[bool] = False,
+    version: str | None = None,
+    force: bool | None = False,
 ):
     client = StudioClient(team=team_name)
     response = client.rm_dataset(name, namespace, project, version, force)
@@ -271,12 +272,21 @@ def save_config(hostname, token, level=ConfigLevel.GLOBAL):
     return config.config_file()
 
 
-def parse_start_time(start_time_str: Optional[str]) -> Optional[str]:
+def parse_start_time(start_time_str: str | None) -> str | None:
     if not start_time_str:
         return None
 
-    # Parse the datetime string using dateparser
-    parsed_datetime = dateparser.parse(start_time_str)
+    # dateparser#1246: it explores strptime patterns lacking a year, which
+    # triggers a CPython 3.13 DeprecationWarning. Suppress that noise until a
+    # new dateparser release includes the upstream fix.
+    # https://github.com/scrapinghub/dateparser/issues/1246
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore",
+            category=DeprecationWarning,
+            module="dateparser\\.utils\\.strptime",
+        )
+        parsed_datetime = dateparser.parse(start_time_str)
 
     if parsed_datetime is None:
         raise DataChainError(
@@ -343,21 +353,21 @@ def show_logs_from_client(client, job_id):
 
 def create_job(
     query_file: str,
-    team_name: Optional[str],
-    env_file: Optional[str] = None,
-    env: Optional[list[str]] = None,
-    workers: Optional[int] = None,
-    files: Optional[list[str]] = None,
-    python_version: Optional[str] = None,
-    repository: Optional[str] = None,
-    req: Optional[list[str]] = None,
-    req_file: Optional[str] = None,
-    priority: Optional[int] = None,
-    cluster: Optional[str] = None,
-    start_time: Optional[str] = None,
-    cron: Optional[str] = None,
-    no_wait: Optional[bool] = False,
-    credentials_name: Optional[str] = None,
+    team_name: str | None,
+    env_file: str | None = None,
+    env: list[str] | None = None,
+    workers: int | None = None,
+    files: list[str] | None = None,
+    python_version: str | None = None,
+    repository: str | None = None,
+    req: list[str] | None = None,
+    req_file: str | None = None,
+    priority: int | None = None,
+    cluster: str | None = None,
+    start_time: str | None = None,
+    cron: str | None = None,
+    no_wait: bool | None = False,
+    credentials_name: str | None = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -403,14 +413,14 @@ def create_job(
     if not response.data:
         raise DataChainError("Failed to create job")
 
-    job_id = response.data.get("job", {}).get("id")
+    job_id = response.data.get("id")
 
     if parsed_start_time or cron:
         print(f"Job {job_id} is scheduled as a task in Studio.")
         return 0
 
     print(f"Job {job_id} created")
-    print("Open the job in Studio at", response.data.get("job", {}).get("url"))
+    print("Open the job in Studio at", response.data.get("url"))
     print("=" * 40)
 
     return 0 if no_wait else show_logs_from_client(client, job_id)
@@ -421,21 +431,19 @@ def upload_files(client: StudioClient, files: list[str]) -> list[str]:
     for file in files:
         file_name = os.path.basename(file)
         with open(file, "rb") as f:
-            file_content = f.read()
-        response = client.upload_file(file_content, file_name)
+            response = client.upload_file(f, file_name)
         if not response.ok:
             raise DataChainError(response.message)
 
         if not response.data:
             raise DataChainError(f"Failed to upload file {file_name}")
 
-        file_id = response.data.get("blob", {}).get("id")
-        if file_id:
+        if file_id := response.data.get("id"):
             file_ids.append(str(file_id))
     return file_ids
 
 
-def cancel_job(job_id: str, team_name: Optional[str]):
+def cancel_job(job_id: str, team_name: str | None):
     token = Config().read().get("studio", {}).get("token")
     if not token:
         raise DataChainError(
@@ -450,13 +458,13 @@ def cancel_job(job_id: str, team_name: Optional[str]):
     print(f"Job {job_id} canceled")
 
 
-def list_jobs(status: Optional[str], team_name: Optional[str], limit: int):
+def list_jobs(status: str | None, team_name: str | None, limit: int):
     client = StudioClient(team=team_name)
     response = client.get_jobs(status, limit)
     if not response.ok:
         raise DataChainError(response.message)
 
-    jobs = response.data.get("jobs", [])
+    jobs = response.data or []
     if not jobs:
         print("No jobs found")
         return
@@ -475,7 +483,7 @@ def list_jobs(status: Optional[str], team_name: Optional[str], limit: int):
     print(tabulate.tabulate(rows, headers="keys", tablefmt="grid"))
 
 
-def show_job_logs(job_id: str, team_name: Optional[str]):
+def show_job_logs(job_id: str, team_name: str | None):
     token = Config().read().get("studio", {}).get("token")
     if not token:
         raise DataChainError(
@@ -486,13 +494,13 @@ def show_job_logs(job_id: str, team_name: Optional[str]):
     return show_logs_from_client(client, job_id)
 
 
-def list_clusters(team_name: Optional[str]):
+def list_clusters(team_name: str | None):
     client = StudioClient(team=team_name)
     response = client.get_clusters()
     if not response.ok:
         raise DataChainError(response.message)
 
-    clusters = response.data.get("clusters", [])
+    clusters = response.data or []
     if not clusters:
         print("No clusters found")
         return
@@ -505,6 +513,7 @@ def list_clusters(team_name: Optional[str]):
             "Cloud Provider": cluster.get("cloud_provider"),
             "Cloud Credentials": cluster.get("cloud_credentials"),
             "Is Active": cluster.get("is_active"),
+            "Is Default": cluster.get("default"),
             "Max Workers": cluster.get("max_workers"),
         }
         for cluster in clusters

datachain/toolkit/split.py

@@ -1,7 +1,7 @@
 import random
-from typing import Optional
 
 from datachain import C, DataChain
+from datachain.lib.signal_schema import SignalResolvingError
 
 RESOLUTION = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
 
@@ -9,7 +9,7 @@ RESOLUTION = 2**31 - 1  # Maximum positive value for a 32-bit signed integer.
 def train_test_split(
     dc: DataChain,
     weights: list[float],
-    seed: Optional[int] = None,
+    seed: int | None = None,
 ) -> list[DataChain]:
     """
     Splits a DataChain into multiple subsets based on the provided weights.
@@ -60,7 +60,10 @@ def train_test_split(
         ```
 
     Note:
-        The splits are random but deterministic, based on Dataset `sys__rand` field.
+        Splits reuse the same best-effort shuffle used by `DataChain.shuffle`. Results
+        are typically repeatable, but earlier operations such as `merge`, `union`, or
+        custom SQL that reshuffle rows can change the outcome between runs. Add order by
+        stable keys first when you need strict reproducibility.
     """
     if len(weights) < 2:
         raise ValueError("Weights should have at least two elements")
@@ -69,16 +72,34 @@ def train_test_split(
 
     weights_normalized = [weight / sum(weights) for weight in weights]
 
+    try:
+        dc.signals_schema.resolve("sys.rand")
+    except SignalResolvingError:
+        dc = dc.persist()
+
     rand_col = C("sys.rand")
     if seed is not None:
         uniform_seed = random.Random(seed).randrange(1, RESOLUTION)  # noqa: S311
         rand_col = (rand_col % RESOLUTION) * uniform_seed  # type: ignore[assignment]
     rand_col = rand_col % RESOLUTION  # type: ignore[assignment]
 
-    return [
-        dc.filter(
-            rand_col >= round(sum(weights_normalized[:index]) * (RESOLUTION - 1)),
-            rand_col < round(sum(weights_normalized[: index + 1]) * (RESOLUTION - 1)),
-        )
-        for index, _ in enumerate(weights_normalized)
-    ]
+    boundaries: list[int] = [0]
+    cumulative = 0.0
+    for weight in weights_normalized[:-1]:
+        cumulative += weight
+        boundary = round(cumulative * RESOLUTION)
+        boundaries.append(min(boundary, RESOLUTION))
+    boundaries.append(RESOLUTION)
+
+    splits: list[DataChain] = []
+    last_index = len(weights_normalized) - 1
+    for index in range(len(weights_normalized)):
+        lower = boundaries[index]
+        if index == last_index:
+            condition = rand_col >= lower
+        else:
+            upper = boundaries[index + 1]
+            condition = (rand_col >= lower) & (rand_col < upper)
+        splits.append(dc.filter(condition))
+
+    return splits