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

datachain/sql/sqlite/base.py

@@ -2,13 +2,12 @@ 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 orjson
 import sqlalchemy as sa
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext.compiler import compiles
@@ -16,6 +15,7 @@ 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,
@@ -88,6 +88,9 @@ def setup():
     compiles(sql_path.file_ext, "sqlite")(compile_path_file_ext)
     compiles(array.length, "sqlite")(compile_array_length)
     compiles(array.contains, "sqlite")(compile_array_contains)
+    compiles(array.slice, "sqlite")(compile_array_slice)
+    compiles(array.join, "sqlite")(compile_array_join)
+    compiles(array.get_element, "sqlite")(compile_array_get_element)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
     compiles(string.regexp_replace, "sqlite")(compile_string_regexp_replace)
@@ -109,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:
@@ -129,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.
@@ -143,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):
@@ -179,7 +196,7 @@ def missing_vector_function(name, exc):
 
 
 def sqlite_string_split(string: str, sep: str, maxsplit: int = -1) -> str:
-    return orjson.dumps(string.split(sep, maxsplit)).decode("utf-8")
+    return json.dumps(string.split(sep, maxsplit), ensure_ascii=False)
 
 
 def sqlite_int_hash_64(x: int) -> int:
@@ -198,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:
@@ -212,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:
@@ -270,6 +285,22 @@ def register_user_defined_sql_functions() -> None:
 
     _registered_function_creators["string_functions"] = create_string_functions
 
+    def create_array_functions(conn):
+        conn.create_function(
+            "json_array_get_element", 2, py_json_array_get_element, deterministic=True
+        )
+        conn.create_function(
+            "json_array_slice", 2, py_json_array_slice, deterministic=True
+        )
+        conn.create_function(
+            "json_array_slice", 3, py_json_array_slice, deterministic=True
+        )
+        conn.create_function(
+            "json_array_join", 2, py_json_array_join, deterministic=True
+        )
+
+    _registered_function_creators["array_functions"] = create_array_functions
+
     has_json_extension = functions_exist(["json_array_length", "json_array_contains"])
     if not has_json_extension:
 
@@ -285,7 +316,11 @@ def register_user_defined_sql_functions() -> None:
 
 
 def adapt_datetime(val: datetime) -> str:
-    if not (val.tzinfo is timezone.utc or val.tzname() == "UTC"):
+    is_utc_check = val.tzinfo is timezone.utc
+    tzname_check = val.tzname() == "UTC"
+    combined_check = is_utc_check or tzname_check
+
+    if not combined_check:
         try:
             val = val.astimezone(timezone.utc)
         except (OverflowError, ValueError, OSError):
@@ -295,6 +330,7 @@ def adapt_datetime(val: datetime) -> str:
                 val = datetime.min.replace(tzinfo=timezone.utc)
             else:
                 raise
+
     return val.replace(tzinfo=None).isoformat(" ")
 
 
@@ -429,13 +465,42 @@ def compile_byte_hamming_distance(element, compiler, **kwargs):
 
 
 def py_json_array_length(arr):
-    return len(orjson.loads(arr))
+    return len(json.loads(arr))
 
 
 def py_json_array_contains(arr, value, is_json):
     if is_json:
-        value = orjson.loads(value)
-    return value in orjson.loads(arr)
+        value = json.loads(value)
+    return value in json.loads(arr)
+
+
+def py_json_array_get_element(val, idx):
+    arr = json.loads(val)
+    try:
+        return arr[idx]
+    except IndexError:
+        return None
+
+
+def py_json_array_slice(val, offset: int, length: int | None = None):
+    arr = json.loads(val)
+    try:
+        return json.dumps(
+            list(arr[offset : offset + length] if length is not None else arr[offset:]),
+            ensure_ascii=False,
+        )
+    except IndexError:
+        return None
+
+
+def py_json_array_join(val, sep: str):
+    return sep.join(json.loads(val))
+
+
+def compile_array_get_element(element, compiler, **kwargs):
+    return compiler.process(
+        func.json_array_get_element(*element.clauses.clauses), **kwargs
+    )
 
 
 def compile_array_length(element, compiler, **kwargs):
@@ -448,6 +513,14 @@ def compile_array_contains(element, compiler, **kwargs):
     )
 
 
+def compile_array_slice(element, compiler, **kwargs):
+    return compiler.process(func.json_array_slice(*element.clauses.clauses), **kwargs)
+
+
+def compile_array_join(element, compiler, **kwargs):
+    return compiler.process(func.json_array_join(*element.clauses.clauses), **kwargs)
+
+
 def compile_string_length(element, compiler, **kwargs):
     return compiler.process(func.length(*element.clauses.clauses), **kwargs)
 
@@ -544,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 orjson
 from sqlalchemy import types
 
+from datachain import json
 from datachain.sql.types import TypeConverter, TypeReadConverter
 
 try:
@@ -28,26 +28,21 @@ class Array(types.UserDefinedType):
 
 
 def adapt_array(arr):
-    return orjson.dumps(arr).decode("utf-8")
+    return json.dumps(arr, ensure_ascii=False)
 
 
 def adapt_dict(dct):
-    return orjson.dumps(dct).decode("utf-8")
+    return json.dumps(dct, ensure_ascii=False)
 
 
 def convert_array(arr):
-    return orjson.loads(arr)
+    return json.loads(arr)
 
 
 def adapt_np_array(arr):
-    def _json_serialize(obj):
-        if isinstance(obj, np.ndarray):
-            return obj.tolist()
-        return obj
-
-    return orjson.dumps(
-        arr, option=orjson.OPT_SERIALIZE_NUMPY, default=_json_serialize
-    ).decode("utf-8")
+    # Primarily needed for UDF numpy results (e.g. WDS)
+    # tolist() gives nested Python lists + native scalars; ujson.dumps handles NaN/Inf.
+    return json.dumps(arr.tolist(), ensure_ascii=False)
 
 
 def adapt_np_generic(val):
@@ -74,5 +69,5 @@ class SQLiteTypeConverter(TypeConverter):
 class SQLiteTypeReadConverter(TypeReadConverter):
     def array(self, value, item_type, dialect):
         if isinstance(value, str):
-            value = orjson.loads(value)
+            value = json.loads(value)
         return super().array(value, item_type, dialect)

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 orjson
 import sqlalchemy as sa
 from sqlalchemy import TypeDecorator, types
 
+from datachain import json as jsonlib
 from datachain.lib.data_model import StandardType
 
 _registry: dict[str, "TypeConverter"] = {}
@@ -58,9 +59,14 @@ def converter(dialect) -> "TypeConverter":
     try:
         return registry[name]
     except KeyError:
-        raise ValueError(
-            f"No type converter registered for dialect: {dialect.name!r}"
-        ) from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No type converter registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def read_converter(dialect) -> "TypeReadConverter":
@@ -68,9 +74,14 @@ def read_converter(dialect) -> "TypeReadConverter":
     try:
         return read_converter_registry[name]
     except KeyError:
-        raise ValueError(
-            f"No read type converter registered for dialect: {dialect.name!r}"
-        ) from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return read_converter_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No read type converter registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def type_defaults(dialect) -> "TypeDefaults":
@@ -78,7 +89,14 @@ def type_defaults(dialect) -> "TypeDefaults":
     try:
         return type_defaults_registry[name]
     except KeyError:
-        raise ValueError(f"No type defaults registered for dialect: {name!r}") from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return type_defaults_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No type defaults registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 def db_defaults(dialect) -> "DBDefaults":
@@ -86,7 +104,14 @@ def db_defaults(dialect) -> "DBDefaults":
     try:
         return db_defaults_registry[name]
     except KeyError:
-        raise ValueError(f"No DB defaults registered for dialect: {name!r}") from None
+        # Fall back to default converter if specific dialect not found
+        try:
+            return db_defaults_registry["default"]
+        except KeyError:
+            raise ValueError(
+                f"No DB defaults registered for dialect: {dialect.name!r} "
+                f"and no default converter available"
+            ) from None
 
 
 class SQLType(TypeDecorator):
@@ -312,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):
@@ -328,7 +371,7 @@ class Array(SQLType):
     def on_read_convert(self, value, dialect):
         r = read_converter(dialect).array(value, self.item_type, dialect)
         if isinstance(self.item_type, JSON):
-            r = [orjson.loads(item) if isinstance(item, str) else item for item in r]
+            r = [jsonlib.loads(item) if isinstance(item, str) else item for item in r]
         return r
 
 
@@ -403,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):
@@ -442,7 +497,7 @@ class TypeReadConverter:
         if isinstance(value, str):
             if value == "":
                 return {}
-            return orjson.loads(value)
+            return jsonlib.loads(value)
         return value
 
     def datetime(self, value):