datachain 0.30.0__py3-none-any.whl → 0.30.2__py3-none-any.whl

datachain/lib/dc/database.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 import sqlalchemy
 
 from datachain.query.schema import ColumnMeta
+from datachain.utils import batched
 
 DEFAULT_DATABASE_BATCH_SIZE = 10_000
 
@@ -74,6 +75,7 @@ def to_database(
     *,
     batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
     on_conflict: Optional[str] = None,
+    conflict_columns: Optional[list[str]] = None,
     column_mapping: Optional[dict[str, Optional[str]]] = None,
 ) -> None:
     """
@@ -82,8 +84,6 @@ def to_database(
     This is the core implementation that handles the actual database operations.
     For user-facing documentation, see DataChain.to_database() method.
     """
-    from datachain.utils import batched
-
     if on_conflict and on_conflict not in ("ignore", "update"):
         raise ValueError(
             f"on_conflict must be 'ignore' or 'update', got: {on_conflict}"
@@ -105,19 +105,26 @@ def to_database(
         metadata = sqlalchemy.MetaData()
         table = sqlalchemy.Table(table_name, metadata, *columns)
 
-        # Check if table already exists to determine if we should clean up on error.
-        inspector = sqlalchemy.inspect(conn)
-        assert inspector  # to satisfy mypy
-        table_existed_before = table_name in inspector.get_table_names()
-
+        table_existed_before = False
         try:
-            table.create(conn, checkfirst=True)
-            rows_iter = chain._leaf_values()
-            for batch in batched(rows_iter, batch_rows):
-                _process_batch(
-                    conn, table, batch, on_conflict, column_indices_and_names
-                )
-            conn.commit()
+            with conn.begin():
+                # Check if table exists to determine if we should clean up on error.
+                inspector = sqlalchemy.inspect(conn)
+                assert inspector  # to satisfy mypy
+                table_existed_before = table_name in inspector.get_table_names()
+
+                table.create(conn, checkfirst=True)
+
+                rows_iter = chain._leaf_values()
+                for batch in batched(rows_iter, batch_rows):
+                    _process_batch(
+                        conn,
+                        table,
+                        batch,
+                        on_conflict,
+                        conflict_columns,
+                        column_indices_and_names,
+                    )
         except Exception:
             if not table_existed_before:
                 try:
@@ -183,7 +190,9 @@ def _prepare_columns(all_columns, column_mapping):
     return column_indices_and_names, columns
 
 
-def _process_batch(conn, table, batch, on_conflict, column_indices_and_names):
+def _process_batch(
+    conn, table, batch, on_conflict, conflict_columns, column_indices_and_names
+):
     """Process a batch of rows with conflict resolution."""
 
     def prepare_row(row_values):
@@ -217,7 +226,19 @@ def _process_batch(conn, table, batch, on_conflict, column_indices_and_names):
             update_values = {
                 col.name: insert_stmt.excluded[col.name] for col in table.columns
             }
-            insert_stmt = insert_stmt.on_conflict_do_update(set_=update_values)
+            if conn.engine.name == "postgresql":
+                if not conflict_columns:
+                    raise ValueError(
+                        "conflict_columns parameter is required when "
+                        "on_conflict='update' with PostgreSQL. Specify the column "
+                        "names that form a unique constraint."
+                    )
+
+                insert_stmt = insert_stmt.on_conflict_do_update(
+                    index_elements=conflict_columns, set_=update_values
+                )
+            else:
+                insert_stmt = insert_stmt.on_conflict_do_update(set_=update_values)
     elif on_conflict:
         import warnings
 

datachain/lib/dc/datachain.py

@@ -2296,6 +2296,7 @@ class DataChain:
         *,
         batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
         on_conflict: Optional[str] = None,
+        conflict_columns: Optional[list[str]] = None,
         column_mapping: Optional[dict[str, Optional[str]]] = None,
     ) -> None:
         """Save chain to a database table using a given database connection.
@@ -2319,6 +2320,9 @@ class DataChain:
                   (default)
                 - "ignore": Skip duplicate rows silently
                 - "update": Update existing rows with new values
+            conflict_columns: List of column names that form a unique constraint
+                for conflict resolution. Required when on_conflict='update' and
+                using PostgreSQL.
             column_mapping: Optional mapping to rename or skip columns:
                 - Dict mapping DataChain column names to database column names
                 - Set values to None to skip columns entirely, or use `defaultdict` to
@@ -2377,6 +2381,7 @@ class DataChain:
             connection,
             batch_rows=batch_rows,
             on_conflict=on_conflict,
+            conflict_columns=conflict_columns,
             column_mapping=column_mapping,
         )
 

datachain/query/dataset.py

@@ -665,7 +665,7 @@ class UDFSignal(UDFStep):
         original_cols = [c for c in subq.c if c.name not in partition_col_names]
 
         # new signal columns that are added to udf_table
-        signal_cols = [c for c in udf_table.c if c.name != "sys__id"]
+        signal_cols = [c for c in udf_table.c if not c.name.startswith("sys__")]
         signal_name_cols = {c.name: c for c in signal_cols}
         cols = signal_cols
 

datachain/sql/__init__.py

@@ -1,6 +1,8 @@
 from sqlalchemy.sql.elements import literal
 from sqlalchemy.sql.expression import column
 
+# Import PostgreSQL dialect registration (registers PostgreSQL type converter)
+from . import postgresql_dialect  # noqa: F401
 from .default import setup as default_setup
 from .selectable import select, values
 

datachain/sql/postgresql_dialect.py

@@ -0,0 +1,9 @@
+"""
+PostgreSQL dialect registration for DataChain.
+"""
+
+from datachain.sql.postgresql_types import PostgreSQLTypeConverter
+from datachain.sql.types import register_backend_types
+
+# Register PostgreSQL type converter
+register_backend_types("postgresql", PostgreSQLTypeConverter())

datachain/sql/postgresql_types.py

@@ -0,0 +1,21 @@
+"""
+PostgreSQL-specific type converter for DataChain.
+
+Handles PostgreSQL-specific type mappings that differ from the default dialect.
+"""
+
+from sqlalchemy.dialects import postgresql
+
+from datachain.sql.types import TypeConverter
+
+
+class PostgreSQLTypeConverter(TypeConverter):
+    """PostgreSQL-specific type converter."""
+
+    def datetime(self):
+        """PostgreSQL uses TIMESTAMP WITH TIME ZONE to preserve timezone information."""
+        return postgresql.TIMESTAMP(timezone=True)
+
+    def json(self):
+        """PostgreSQL uses JSONB for better performance and query capabilities."""
+        return postgresql.JSONB()

datachain/sql/sqlite/__init__.py

@@ -1,4 +1,8 @@
-from .base import create_user_defined_sql_functions, setup, sqlite_dialect
+from .base import (
+    create_user_defined_sql_functions,
+    setup,
+    sqlite_dialect,
+)
 
 __all__ = [
     "create_user_defined_sql_functions",

datachain/sql/sqlite/base.py

@@ -304,7 +304,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):
@@ -314,6 +318,7 @@ def adapt_datetime(val: datetime) -> str:
                 val = datetime.min.replace(tzinfo=timezone.utc)
             else:
                 raise
+
     return val.replace(tzinfo=None).isoformat(" ")
 
 

datachain/sql/types.py

@@ -58,9 +58,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 +73,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 +88,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 +103,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):

{datachain-0.30.0.dist-info → datachain-0.30.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.30.0
+Version: 0.30.2
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License-Expression: Apache-2.0
@@ -81,8 +81,10 @@ Provides-Extra: video
 Requires-Dist: ffmpeg-python; extra == "video"
 Requires-Dist: imageio[ffmpeg,pyav]>=2.37.0; extra == "video"
 Requires-Dist: opencv-python; extra == "video"
+Provides-Extra: postgres
+Requires-Dist: psycopg2-binary>=2.9.0; extra == "postgres"
 Provides-Extra: tests
-Requires-Dist: datachain[audio,hf,remote,torch,vector,video]; extra == "tests"
+Requires-Dist: datachain[audio,hf,postgres,remote,torch,vector,video]; extra == "tests"
 Requires-Dist: pytest<9,>=8; extra == "tests"
 Requires-Dist: pytest-sugar>=0.9.6; extra == "tests"
 Requires-Dist: pytest-cov>=4.1.0; extra == "tests"
@@ -90,6 +92,7 @@ Requires-Dist: pytest-mock>=3.12.0; extra == "tests"
 Requires-Dist: pytest-servers[all]>=0.5.9; extra == "tests"
 Requires-Dist: pytest-benchmark[histogram]; extra == "tests"
 Requires-Dist: pytest-xdist>=3.3.1; extra == "tests"
+Requires-Dist: pytest-env>=1.1.0; extra == "tests"
 Requires-Dist: virtualenv; extra == "tests"
 Requires-Dist: dulwich; extra == "tests"
 Requires-Dist: hypothesis; extra == "tests"

{datachain-0.30.0.dist-info → datachain-0.30.2.dist-info}/RECORD

@@ -103,8 +103,8 @@ datachain/lib/convert/unflatten.py,sha256=ysMkstwJzPMWUlnxn-Z-tXJR3wmhjHeSN_P-sD
 datachain/lib/convert/values_to_tuples.py,sha256=j5yZMrVUH6W7b-7yUvdCTGI7JCUAYUOzHUGPoyZXAB0,4360
 datachain/lib/dc/__init__.py,sha256=TFci5HTvYGjBesNUxDAnXaX36PnzPEUSn5a6JxB9o0U,872
 datachain/lib/dc/csv.py,sha256=q6a9BpapGwP6nwy6c5cklxQumep2fUp9l2LAjtTJr6s,4411
-datachain/lib/dc/database.py,sha256=MPE-KzwcR2DhWLCEbl1gWFp63dLqjWuiJ1iEfC2BrJI,12443
-datachain/lib/dc/datachain.py,sha256=_C9PZjUHVewpdp94AR2GS3QEI96Svsyx52dLJVM4tm4,98143
+datachain/lib/dc/database.py,sha256=4Fhen6KZRMYzSfONydwTFwSUECbdff0t-9GSj6ADyYM,13288
+datachain/lib/dc/datachain.py,sha256=UHICzncqG6GmDXxrX0DEYVXJK19c-8H-eoRuA7097zc,98439
 datachain/lib/dc/datasets.py,sha256=P6CIJizD2IYFwOQG5D3VbQRjDmUiRH0ysdtb551Xdm8,15098
 datachain/lib/dc/hf.py,sha256=AP_MUHg6HJWae10PN9hD_beQVjrl0cleZ6Cvhtl1yoI,2901
 datachain/lib/dc/json.py,sha256=dNijfJ-H92vU3soyR7X1IiDrWhm6yZIGG3bSnZkPdAE,2733
@@ -126,7 +126,7 @@ datachain/model/ultralytics/pose.py,sha256=pBlmt63Qe68FKmexHimUGlNbNOoOlMHXG4fzX
 datachain/model/ultralytics/segment.py,sha256=63bDCj43E6iZ0hFI5J6uQfksdCmjEp6sEm1XzVaE8pw,2986
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
 datachain/query/batch.py,sha256=-goxLpE0EUvaDHu66rstj53UnfHpYfBUGux8GSpJ93k,4306
-datachain/query/dataset.py,sha256=OJZ_YwpS5i4B0wVmosMmMNW1qABr6zyOmqNHQdAWir4,62704
+datachain/query/dataset.py,sha256=hUKKHuqkfNaojzOt6rMIksU-PG72i_lfbKSHZ5rt--M,62715
 datachain/query/dispatch.py,sha256=A0nPxn6mEN5d9dDo6S8m16Ji_9IvJLXrgF2kqXdi4fs,15546
 datachain/query/metrics.py,sha256=DOK5HdNVaRugYPjl8qnBONvTkwjMloLqAr7Mi3TjCO0,858
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
@@ -137,9 +137,11 @@ datachain/query/udf.py,sha256=e753bDJzTNjGFQn1WGTvOAWSwjDbrFI1-_DDWkWN2ls,1343
 datachain/query/utils.py,sha256=a2PTBZ3qsG6XlUcp9XsoGiQfKkca4Q3m-VzFgiGQPAc,1230
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/remote/studio.py,sha256=pDThxvEEpIKVGfa9rmtz_zeqHwrgzh0Lv-Pd4wzDx5k,15448
-datachain/sql/__init__.py,sha256=6SQRdbljO3d2hx3EAVXEZrHQKv5jth0Jh98PogT59No,262
+datachain/sql/__init__.py,sha256=8D2omsBiATt8bjLjGo6jBEtaKEkOlnlNFWhVryHMDv0,388
+datachain/sql/postgresql_dialect.py,sha256=pDTfH8xaXz5xZsq8O1aQUvWLRIv_ogYeAqtmKlPp3Rw,280
+datachain/sql/postgresql_types.py,sha256=ryb_0lzuA9UOJ_B6nW9Yb8nJjzeSmEItAL_Ceue65lc,627
 datachain/sql/selectable.py,sha256=cTc60qVoAwqqss0Vop8Lt5Z-ROnM1XrQmL_GLjRxhXs,1765
-datachain/sql/types.py,sha256=ASSPkmM5EzdRindqj2O7WHLXq8VHAgFYedG8lYfGvVI,14045
+datachain/sql/types.py,sha256=RWOghtYFx14K-e71QOGg5yfKb-A4-4JgFjaJ0wCZ17Y,15006
 datachain/sql/utils.py,sha256=rzlJw08etivdrcuQPqNVvVWhuVSyUPUQEEc6DOhu258,818
 datachain/sql/default/__init__.py,sha256=XQ2cEZpzWiABqjV-6yYHUBGI9vN_UHxbxZENESmVAWw,45
 datachain/sql/default/base.py,sha256=QD-31C6JnyOXzogyDx90sUhm7QvgXIYpeHEASH84igU,628
@@ -151,16 +153,16 @@ datachain/sql/functions/numeric.py,sha256=BK2KCiPSgM2IveCq-9M_PG3CtPBlztaS9TTn1L
 datachain/sql/functions/path.py,sha256=zixpERotTFP6LZ7I4TiGtyRA8kXOoZmH1yzH9oRW0mg,1294
 datachain/sql/functions/random.py,sha256=vBwEEj98VH4LjWixUCygQ5Bz1mv1nohsCG0-ZTELlVg,271
 datachain/sql/functions/string.py,sha256=E-T9OIzUR-GKaLgjZsEtg5CJrY_sLf1lt1awTvY7w2w,1426
-datachain/sql/sqlite/__init__.py,sha256=TAdJX0Bg28XdqPO-QwUVKy8rg78cgMileHvMNot7d04,166
-datachain/sql/sqlite/base.py,sha256=ldyWoMBvrYH789GK8xE4wbBNYxz4B-g4BmQAK6_KjR0,21398
+datachain/sql/sqlite/__init__.py,sha256=PsLaDSij9a03VxGSpagpNl7NQsGtgm72ArUeALZONoc,183
+datachain/sql/sqlite/base.py,sha256=6aoQHeggY3hs31_YZ-wlYKA1Lto4MFOpgfgRspH6IMc,21498
 datachain/sql/sqlite/types.py,sha256=cH6oge2E_YWFy22wY-txPJH8gxoQFSpCthtZR8PZjpo,1849
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.30.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.30.0.dist-info/METADATA,sha256=khld2YGTkc6UpViqw2RxeeDeUwG8wGW6X2RR2zBaydk,13766
-datachain-0.30.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.30.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.30.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.30.0.dist-info/RECORD,,
+datachain-0.30.2.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.30.2.dist-info/METADATA,sha256=nLnOc_mmzRuLWdFk3hiUi_P71TLY975X_ZWk4iyojeg,13910
+datachain-0.30.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.30.2.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.30.2.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.30.2.dist-info/RECORD,,