datachain 0.3.4__py3-none-any.whl → 0.3.5__py3-none-any.whl

datachain/lib/convert/flatten.py

@@ -1,34 +1,6 @@
-from datetime import datetime
-
 from pydantic import BaseModel
 
 from datachain.lib.model_store import ModelStore
-from datachain.sql.types import (
-    JSON,
-    Array,
-    Binary,
-    Boolean,
-    DateTime,
-    Float,
-    Int,
-    Int32,
-    Int64,
-    NullType,
-    String,
-)
-
-DATACHAIN_TO_TYPE = {
-    Int: int,
-    Int32: int,
-    Int64: int,
-    String: str,
-    Float: float,
-    Boolean: bool,
-    DateTime: datetime,
-    Binary: bytes,
-    Array(NullType): list,
-    JSON: dict,
-}
 
 
 def flatten(obj: BaseModel):

datachain/lib/dc.py

@@ -1224,14 +1224,11 @@ class DataChain(DatasetQuery):
         """
         headers, max_length = self._effective_signals_schema.get_headers_with_length()
         if flatten or max_length < 2:
-            columns = []
-            if headers:
-                columns = [".".join(filter(None, header)) for header in headers]
-            return pd.DataFrame.from_records(self.to_records(), columns=columns)
+            columns = [".".join(filter(None, header)) for header in headers]
+        else:
+            columns = pd.MultiIndex.from_tuples(map(tuple, headers))
 
-        return pd.DataFrame(
-            self.results(), columns=pd.MultiIndex.from_tuples(map(tuple, headers))
-        )
+        return pd.DataFrame.from_records(self.results(), columns=columns)
 
     def show(
         self,
@@ -1524,6 +1521,7 @@ class DataChain(DatasetQuery):
         to_insert: Optional[Union[dict, list[dict]]],
         session: Optional[Session] = None,
         in_memory: bool = False,
+        schema: Optional[dict[str, DataType]] = None,
     ) -> "DataChain":
         """Create a DataChain from the provided records. This method can be used for
         programmatically generating a chain in contrast of reading data from storages
@@ -1532,10 +1530,10 @@ class DataChain(DatasetQuery):
         Parameters:
             to_insert : records (or a single record) to insert. Each record is
                         a dictionary of signals and theirs values.
+            schema : describes chain signals and their corresponding types
 
         Example:
             ```py
-            empty = DataChain.from_records()
             single_record = DataChain.from_records(DataChain.DEFAULT_FILE_RECORD)
             ```
         """
@@ -1543,11 +1541,27 @@ class DataChain(DatasetQuery):
         catalog = session.catalog
 
         name = session.generate_temp_dataset_name()
-        columns: tuple[sqlalchemy.Column[Any], ...] = tuple(
-            sqlalchemy.Column(name, typ)
-            for name, typ in File._datachain_column_types.items()
+        signal_schema = None
+        columns: list[sqlalchemy.Column] = []
+
+        if schema:
+            signal_schema = SignalSchema(schema)
+            columns = signal_schema.db_signals(as_columns=True)  # type: ignore[assignment]
+        else:
+            columns = [
+                sqlalchemy.Column(name, typ)
+                for name, typ in File._datachain_column_types.items()
+            ]
+
+        dsr = catalog.create_dataset(
+            name,
+            columns=columns,
+            feature_schema=(
+                signal_schema.clone_without_sys_signals().serialize()
+                if signal_schema
+                else None
+            ),
         )
-        dsr = catalog.create_dataset(name, columns=columns)
 
         if isinstance(to_insert, dict):
             to_insert = [to_insert]

datachain/lib/signal_schema.py

@@ -2,6 +2,7 @@ import copy
 from collections.abc import Iterator, Sequence
 from dataclasses import dataclass
 from datetime import datetime
+from inspect import isclass
 from typing import (
     TYPE_CHECKING,
     Annotated,
@@ -14,10 +15,10 @@ from typing import (
     get_origin,
 )
 
+import sqlalchemy as sa
 from pydantic import BaseModel, create_model
 from typing_extensions import Literal as LiteralEx
 
-from datachain.lib.convert.flatten import DATACHAIN_TO_TYPE
 from datachain.lib.convert.python_to_sql import python_to_sql
 from datachain.lib.convert.sql_to_python import sql_to_python
 from datachain.lib.convert.unflatten import unflatten_to_json_pos
@@ -26,6 +27,7 @@ from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
 from datachain.query.schema import DEFAULT_DELIMITER, Column
+from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
@@ -104,12 +106,15 @@ class SignalSchema:
     def from_column_types(col_types: dict[str, Any]) -> "SignalSchema":
         signals: dict[str, DataType] = {}
         for field, col_type in col_types.items():
-            if (py_type := DATACHAIN_TO_TYPE.get(col_type, None)) is None:
+            if isinstance(col_type, SQLType):
+                signals[field] = col_type.python_type
+            elif isclass(col_type) and issubclass(col_type, SQLType):
+                signals[field] = col_type().python_type
+            else:
                 raise SignalSchemaError(
                     f"signal schema cannot be obtained for column '{field}':"
-                    f" unsupported type '{py_type}'"
+                    f" unsupported type '{col_type}'"
                 )
-            signals[field] = py_type
         return SignalSchema(signals)
 
     def serialize(self) -> dict[str, str]:
@@ -232,7 +237,7 @@ class SignalSchema:
         signals = [
             DEFAULT_DELIMITER.join(path)
             if not as_columns
-            else Column(DEFAULT_DELIMITER.join(path), python_to_sql(_type))
+            else sa.Column(DEFAULT_DELIMITER.join(path), python_to_sql(_type))
             for path, _type, has_subtree, _ in self.get_flat_tree()
             if not has_subtree
         ]

datachain/query/dataset.py

@@ -878,17 +878,14 @@ class SQLUnion(Step):
         temp_tables.extend(self.query1.temp_table_names)
         q2 = self.query2.apply_steps().select().subquery()
         temp_tables.extend(self.query2.temp_table_names)
-        columns1, columns2 = fill_columns(q1.columns, q2.columns)
+
+        columns1, columns2 = _order_columns(q1.columns, q2.columns)
 
         def q(*columns):
             names = {c.name for c in columns}
             col1 = [c for c in columns1 if c.name in names]
             col2 = [c for c in columns2 if c.name in names]
-            res = (
-                sqlalchemy.select(*col1)
-                .select_from(q1)
-                .union_all(sqlalchemy.select(*col2).select_from(q2))
-            )
+            res = sqlalchemy.select(*col1).union_all(sqlalchemy.select(*col2))
 
             subquery = res.subquery()
             return sqlalchemy.select(*subquery.c).select_from(subquery)
@@ -1021,23 +1018,46 @@ class GroupBy(Step):
         return step_result(q, grouped_query.selected_columns)
 
 
-def fill_columns(
-    *column_iterables: Iterable[ColumnElement],
+def _validate_columns(
+    left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
+) -> set[str]:
+    left_names = {c.name for c in left_columns}
+    right_names = {c.name for c in right_columns}
+
+    if left_names == right_names:
+        return left_names
+
+    missing_right = left_names - right_names
+    missing_left = right_names - left_names
+
+    def _prepare_msg_part(missing_columns: set[str], side: str) -> str:
+        return f"{', '.join(sorted(missing_columns))} only present in {side}"
+
+    msg_parts = [
+        _prepare_msg_part(missing_columns, found_side)
+        for missing_columns, found_side in zip(
+            [
+                missing_right,
+                missing_left,
+            ],
+            ["left", "right"],
+        )
+        if missing_columns
+    ]
+    msg = f"Cannot perform union. {'. '.join(msg_parts)}"
+
+    raise ValueError(msg)
+
+
+def _order_columns(
+    left_columns: Iterable[ColumnElement], right_columns: Iterable[ColumnElement]
 ) -> list[list[ColumnElement]]:
-    column_dicts = [{c.name: c for c in columns} for columns in column_iterables]
-    combined_columns = {n: c for col_dict in column_dicts for n, c in col_dict.items()}
-
-    result: list[list[ColumnElement]] = [[] for _ in column_dicts]
-    for n in combined_columns:
-        col = next(col_dict[n] for col_dict in column_dicts if n in col_dict)
-        for col_dict, out in zip(column_dicts, result):
-            if n in col_dict:
-                out.append(col_dict[n])
-            else:
-                # Cast the NULL to ensure all columns are aware of their type
-                # Label it to ensure it's aware of its name
-                out.append(sqlalchemy.cast(sqlalchemy.null(), col.type).label(n))
-    return result
+    column_order = _validate_columns(left_columns, right_columns)
+    column_dicts = [
+        {c.name: c for c in columns} for columns in [left_columns, right_columns]
+    ]
+
+    return [[d[n] for n in column_order] for d in column_dicts]
 
 
 @attrs.define

datachain/sql/types.py

@@ -20,6 +20,8 @@ from typing import Any, Union
 import sqlalchemy as sa
 from sqlalchemy import TypeDecorator, types
 
+from datachain.lib.data_model import StandardType
+
 _registry: dict[str, "TypeConverter"] = {}
 registry = MappingProxyType(_registry)
 
@@ -91,6 +93,10 @@ class SQLType(TypeDecorator):
     impl: type[types.TypeEngine[Any]] = types.TypeEngine
     cache_ok = True
 
+    @property
+    def python_type(self) -> StandardType:
+        raise NotImplementedError
+
     def to_dict(self) -> dict[str, Any]:
         return {"type": self.__class__.__name__}
 
@@ -103,7 +109,7 @@ class String(SQLType):
     impl = types.String
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return str
 
     def load_dialect_impl(self, dialect):
@@ -125,7 +131,7 @@ class Boolean(SQLType):
     impl = types.Boolean
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return bool
 
     def load_dialect_impl(self, dialect):
@@ -147,7 +153,7 @@ class Int(SQLType):
     impl = types.INTEGER
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return int
 
     def load_dialect_impl(self, dialect):
@@ -217,7 +223,7 @@ class Float(SQLType):
     impl = types.FLOAT
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return float
 
     def load_dialect_impl(self, dialect):
@@ -271,7 +277,7 @@ class Array(SQLType):
     impl = types.ARRAY
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return list
 
     def load_dialect_impl(self, dialect):
@@ -314,7 +320,7 @@ class JSON(SQLType):
     impl = types.JSON
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return dict
 
     def load_dialect_impl(self, dialect):
@@ -336,7 +342,7 @@ class DateTime(SQLType):
     impl = types.DATETIME
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return datetime
 
     def load_dialect_impl(self, dialect):
@@ -358,7 +364,7 @@ class Binary(SQLType):
     impl = types.BINARY
 
     @property
-    def python_type(self):
+    def python_type(self) -> StandardType:
         return bytes
 
     def load_dialect_impl(self, dialect):

{datachain-0.3.4.dist-info → datachain-0.3.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.3.4
+Version: 0.3.5
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -40,6 +40,7 @@ Requires-Dist: pydantic <3,>=2
 Requires-Dist: jmespath >=1.0
 Requires-Dist: datamodel-code-generator >=0.25
 Requires-Dist: Pillow <11,>=10.0.0
+Requires-Dist: msgpack <2,>=1.0.4
 Requires-Dist: numpy <2,>=1 ; sys_platform == "win32"
 Provides-Extra: dev
 Requires-Dist: datachain[docs,tests] ; extra == 'dev'
@@ -61,12 +62,10 @@ Requires-Dist: numpy <2,>=1 ; extra == 'examples'
 Requires-Dist: defusedxml ; extra == 'examples'
 Requires-Dist: accelerate ; extra == 'examples'
 Requires-Dist: unstructured[pdf] ; extra == 'examples'
-Requires-Dist: pdfplumber ==0.11.3 ; extra == 'examples'
+Requires-Dist: pdfplumber ==0.11.4 ; extra == 'examples'
 Requires-Dist: huggingface-hub[hf_transfer] ; extra == 'examples'
-Requires-Dist: nltk ==3.8.1 ; extra == 'examples'
 Provides-Extra: remote
 Requires-Dist: lz4 ; extra == 'remote'
-Requires-Dist: msgpack <2,>=1.0.4 ; extra == 'remote'
 Requires-Dist: requests >=2.22.0 ; extra == 'remote'
 Provides-Extra: tests
 Requires-Dist: datachain[remote,torch,vector] ; extra == 'tests'

{datachain-0.3.4.dist-info → datachain-0.3.5.dist-info}/RECORD

@@ -42,7 +42,7 @@ datachain/lib/arrow.py,sha256=D8N7zCppRdc5sTYT1hNIbROc-sKA_8FN5J_m-KjD3Us,4929
 datachain/lib/clip.py,sha256=16u4b_y2Y15nUS2UN_8ximMo6r_-_4IQpmct2ol-e-g,5730
 datachain/lib/data_model.py,sha256=ZvtMRMcPpBxI-rOhkXb-ry1PkGYcEFFK1w1wH12vs4g,1718
 datachain/lib/dataset_info.py,sha256=lONGr71ozo1DS4CQEhnpKORaU4qFb6Ketv8Xm8CVm2U,2188
-datachain/lib/dc.py,sha256=0pwNb91GW8MnHLfFd2YvEtEH0n77c3nxp5ozwIyW86o,58827
+datachain/lib/dc.py,sha256=atGpaeCUwxDEgHIFmWqG1rAnqe7utT6S7c1jM5yVb7c,59246
 datachain/lib/file.py,sha256=ZHpdilDPYCob8uqtwUPtBvBNxVvQRq4AC_0IGg5m-G4,12003
 datachain/lib/image.py,sha256=TgYhRhzd4nkytfFMeykQkPyzqb5Le_-tU81unVMPn4Q,2328
 datachain/lib/listing.py,sha256=nXLmGae_oQke4hnurzzWiHTEjHjWiqqHdB41Wb-hMTk,3521
@@ -50,7 +50,7 @@ datachain/lib/meta_formats.py,sha256=Hels85LJmNCz1aYVJvhymNdAt3qdJ2-qoxsIiUezrow
 datachain/lib/model_store.py,sha256=c4USXsBBjrGH8VOh4seIgOiav-qHOwdoixtxfLgU63c,2409
 datachain/lib/pytorch.py,sha256=9PsypKseyKfIimTmTQOgb-pbNXgeeAHLdlWx0qRPULY,5660
 datachain/lib/settings.py,sha256=39thOpYJw-zPirzeNO6pmRC2vPrQvt4eBsw1xLWDFsw,2344
-datachain/lib/signal_schema.py,sha256=VL9TR0CJ3eRzjIDr-8e-e7cZKuMBbPUZtY2lGAsucc0,15734
+datachain/lib/signal_schema.py,sha256=MS8qkOIl-3Qh3KyYTCtuSgF9nP5PeaGccbtGqfWo2wI,15902
 datachain/lib/text.py,sha256=dVe2Ilc_gW2EV0kun0UwegiCkapWcd20cef7CgINWHU,1083
 datachain/lib/udf.py,sha256=n3x6No-7l5LAciPJPWwZbA8WtTnGUU7d0wRL6CyfZh8,11847
 datachain/lib/udf_signature.py,sha256=gMStcEeYJka5M6cg50Z9orC6y6HzCAJ3MkFqqn1fjZg,7137
@@ -59,7 +59,7 @@ datachain/lib/vfile.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/webdataset.py,sha256=SsjCKLSKEkHRRfeTHQhjoGqNPqIWw_SCWQcUwgUWWP0,8282
 datachain/lib/webdataset_laion.py,sha256=PQP6tQmUP7Xu9fPuAGK1JDBYA6T5UufYMUTGaxgspJA,2118
 datachain/lib/convert/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/lib/convert/flatten.py,sha256=YMoC00BqEy3zSpvCp6Q0DfxihuPmgjUJj1g2cesWGPs,1790
+datachain/lib/convert/flatten.py,sha256=Uebc5CeqCsacp-nr6IG9i6OGuUavXqdqnoGctZBk3RQ,1384
 datachain/lib/convert/python_to_sql.py,sha256=4gplGlr_Kg-Z40OpJUzJiarDWj7pwbUOk-dPOYYCJ9Q,2629
 datachain/lib/convert/sql_to_python.py,sha256=lGnKzSF_tz9Y_5SSKkrIU95QEjpcDzvOxIRkEKTQag0,443
 datachain/lib/convert/unflatten.py,sha256=Ogvh_5wg2f38_At_1lN0D_e2uZOOpYEvwvB2xdq56Tw,2012
@@ -67,7 +67,7 @@ datachain/lib/convert/values_to_tuples.py,sha256=YOdbjzHq-uj6-cV2Qq43G72eN2avMND
 datachain/query/__init__.py,sha256=tv-spkjUCYamMN9ys_90scYrZ8kJ7C7d1MTYVmxGtk4,325
 datachain/query/batch.py,sha256=-vlpINJiertlnaoUVv1C95RatU0F6zuhpIYRufJRo1M,3660
 datachain/query/builtins.py,sha256=EmKPYsoQ46zwdyOn54MuCzvYFmfsBn5F8zyF7UBUfrc,2550
-datachain/query/dataset.py,sha256=7lxlybS7I5IPsgOqMz-W4vS6kWBDHkHQRqBHlIRYRPw,60473
+datachain/query/dataset.py,sha256=4F_Q101Lbpc0YxOAcP3rc3GtKv8HwxpqF9lpJ0OoUEk,60818
 datachain/query/dispatch.py,sha256=GBh3EZHDp5AaXxrjOpfrpfsuy7Umnqxu-MAXcK9X3gc,12945
 datachain/query/metrics.py,sha256=vsECqbZfoSDBnvC3GQlziKXmISVYDLgHP1fMPEOtKyo,640
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
@@ -79,7 +79,7 @@ datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 datachain/remote/studio.py,sha256=f5s6qSZ9uB4URGUoU_8_W1KZRRQQVSm6cgEBkBUEfuE,7226
 datachain/sql/__init__.py,sha256=A2djrbQwSMUZZEIKGnm-mnRA-NDSbiDJNpAmmwGNyIo,303
 datachain/sql/selectable.py,sha256=fBM-wS1TUA42kVEAAiwqGtibIevyZAEritwt8PZGyLQ,1589
-datachain/sql/types.py,sha256=1MFvECB_5A6QwQKKY3VPhvitgKDlc2aB7iBjY4hv1_s,13034
+datachain/sql/types.py,sha256=1ofJjgzKTxFLl1WaMSI9pLvdHGZ1U24I0z5i-gChqDI,13305
 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
@@ -94,9 +94,9 @@ datachain/sql/sqlite/base.py,sha256=WLPHBhZbXbiqPoRV1VgDrXJqku4UuvJpBhYeQ0k5rI8,
 datachain/sql/sqlite/types.py,sha256=yzvp0sXSEoEYXs6zaYC_2YubarQoZH-MiUNXcpuEP4s,1573
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.3.4.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.3.4.dist-info/METADATA,sha256=nV1-yJcDxoWuaM8uSwEzYCpDqSMxhxPle3EZZ98a-LA,16789
-datachain-0.3.4.dist-info/WHEEL,sha256=nCVcAvsfA9TDtwGwhYaRrlPhTLV9m-Ga6mdyDtuwK18,91
-datachain-0.3.4.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.3.4.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.3.4.dist-info/RECORD,,
+datachain-0.3.5.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.3.5.dist-info/METADATA,sha256=SaQj0C0_Ugll_S1RTRCkFM4U1fZwC7bweiaQZhovqcs,16719
+datachain-0.3.5.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
+datachain-0.3.5.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.3.5.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.3.5.dist-info/RECORD,,

{datachain-0.3.4.dist-info → datachain-0.3.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.0)
+Generator: setuptools (73.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any