pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/io/hf_datasets.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import typing
-from typing import Any, Optional, Union
+from typing import Any
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
@@ -31,12 +31,13 @@ _hf_to_pxt: dict[str, ts.ColumnType] = {
     'timestamp[s]': ts.TimestampType(nullable=True),
     'timestamp[ms]': ts.TimestampType(nullable=True),  # HF dataset iterator converts timestamps to datetime.datetime
     'timestamp[us]': ts.TimestampType(nullable=True),
+    'timestamp[ns]': ts.TimestampType(nullable=True),
     'date32': ts.DateType(nullable=True),
     'date64': ts.DateType(nullable=True),
 }
 
 
-def _to_pixeltable_type(feature_type: Any, nullable: bool) -> Optional[ts.ColumnType]:
+def _to_pixeltable_type(feature_type: Any, nullable: bool) -> ts.ColumnType | None:
     """Convert a huggingface feature type to a pixeltable ColumnType if one is defined."""
     import datasets
 
@@ -47,28 +48,51 @@ def _to_pixeltable_type(feature_type: Any, nullable: bool) -> Optional[ts.Column
         # example: Value(dtype='int64', id=None)
         pt = _hf_to_pxt.get(feature_type.dtype, None)
         return pt.copy(nullable=nullable) if pt is not None else None
-    elif isinstance(feature_type, datasets.Sequence):
+    elif isinstance(feature_type, (datasets.Sequence, datasets.LargeList)):
         # example: cohere wiki. Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None)
         dtype = _to_pixeltable_type(feature_type.feature, nullable)
-        length = feature_type.length if feature_type.length != -1 else None
-        return ts.ArrayType(shape=(length,), dtype=dtype)
+        if dtype is None:
+            return None
+        if dtype.is_int_type() or dtype.is_float_type() or dtype.is_bool_type() or dtype.is_string_type():
+            length = feature_type.length if feature_type.length != -1 else None
+            return ts.ArrayType(shape=(length,), dtype=dtype, nullable=nullable)
+        else:
+            # Sequence of dicts must be cast as Json
+            return ts.JsonType(nullable=nullable)
     elif isinstance(feature_type, datasets.Image):
         return ts.ImageType(nullable=nullable)
+    elif isinstance(feature_type, datasets.Audio):
+        return ts.AudioType(nullable=nullable)
+    elif isinstance(feature_type, datasets.Video):
+        return ts.VideoType(nullable=nullable)
+    elif isinstance(feature_type, (datasets.Array2D, datasets.Array3D, datasets.Array4D, datasets.Array5D)):
+        # Multi-dimensional arrays with fixed shape and dtype
+        inner_dtype = _hf_to_pxt.get(feature_type.dtype, None)
+        if inner_dtype is None:
+            return None
+        return ts.ArrayType(shape=feature_type.shape, dtype=inner_dtype, nullable=nullable)
+    elif isinstance(feature_type, (datasets.Translation, datasets.TranslationVariableLanguages)):
+        # Translation types are dict-like structures
+        return ts.JsonType(nullable=nullable)
+    elif isinstance(feature_type, (list, dict)):
+        return ts.JsonType(nullable=nullable)
     else:
         return None
 
 
-def _get_hf_schema(dataset: Union[datasets.Dataset, datasets.DatasetDict]) -> datasets.Features:
+def _get_hf_schema(dataset: datasets.Dataset | datasets.DatasetDict) -> datasets.Features:
     """Get the schema of a huggingface dataset as a dictionary."""
     import datasets
 
-    first_dataset = dataset if isinstance(dataset, datasets.Dataset) else next(iter(dataset.values()))
+    first_dataset = (
+        dataset if isinstance(dataset, (datasets.Dataset, datasets.IterableDataset)) else next(iter(dataset.values()))
+    )
     return first_dataset.features
 
 
 def huggingface_schema_to_pxt_schema(
     hf_schema: datasets.Features, schema_overrides: dict[str, Any], primary_key: list[str]
-) -> dict[str, Optional[ts.ColumnType]]:
+) -> dict[str, ts.ColumnType | None]:
     """Generate a pixeltable schema from a huggingface dataset schema.
     Columns without a known mapping are mapped to None
     """
@@ -83,20 +107,41 @@ def huggingface_schema_to_pxt_schema(
 
 def import_huggingface_dataset(
     table_path: str,
-    dataset: Union[datasets.Dataset, datasets.DatasetDict],
+    dataset: datasets.Dataset | datasets.DatasetDict | datasets.IterableDataset | datasets.IterableDatasetDict,
     *,
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     **kwargs: Any,
 ) -> pxt.Table:
-    """Create a new base table from a Huggingface dataset, or dataset dict with multiple splits.
-        Requires `datasets` library to be installed.
+    """
+    Create a new base table from a Huggingface dataset, or dataset dict with multiple splits.
+    Requires `datasets` library to be installed.
+
+    HuggingFace feature types are mapped to Pixeltable column types as follows:
+
+    - `Value(bool)`: `Bool`<br/>
+      `Value(int*/uint*)`: `Int`<br/>
+      `Value(float*)`: `Float`<br/>
+      `Value(string/large_string)`: `String`<br/>
+      `Value(timestamp*)`: `Timestamp`<br/>
+      `Value(date*)`: `Date`
+    - `ClassLabel`: `String` (converted to label names)
+    - `Sequence`/`LargeList` of numeric types: `Array`
+    - `Sequence`/`LargeList` of string: `Json`
+    - `Sequence`/`LargeList` of dicts: `Json`
+    - `Array2D`-`Array5D`: `Array` (preserves shape)
+    - `Image`: `Image`
+    - `Audio`: `Audio`
+    - `Video`: `Video`
+    - `Translation`/`TranslationVariableLanguages`: `Json`
 
     Args:
         table_path: Path to the table.
-        dataset: Huggingface [`datasets.Dataset`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.Dataset)
-            or [`datasets.DatasetDict`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.DatasetDict)
-            to insert into the table.
+        dataset: An instance of any of the Huggingface dataset classes:
+            [`datasets.Dataset`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.Dataset),
+            [`datasets.DatasetDict`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.DatasetDict),
+            [`datasets.IterableDataset`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.IterableDataset),
+            [`datasets.IterableDatasetDict`](https://huggingface.co/docs/datasets/en/package_reference/main_classes#datasets.IterableDatasetDict)
         schema_overrides: If specified, then for each (name, type) pair in `schema_overrides`, the column with
             name `name` will be given type `type`, instead of being inferred from the `Dataset` or `DatasetDict`.
             The keys in `schema_overrides` should be the column names of the `Dataset` or `DatasetDict` (whether or not

pixeltable/io/label_studio.py

@@ -4,19 +4,22 @@ import logging
 import os
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Iterator, Literal, Optional, cast
+from typing import Any, Iterator, Literal
 from xml.etree import ElementTree as ET
 
-import label_studio_sdk  # type: ignore[import-untyped]
+import label_studio_sdk
 import PIL.Image
 from requests.exceptions import HTTPError
 
 import pixeltable.type_system as ts
 from pixeltable import Column, Table, env, exceptions as excs
+from pixeltable.catalog import ColumnHandle
+from pixeltable.catalog.update_status import RowCountStats, UpdateStatus
 from pixeltable.config import Config
 from pixeltable.exprs import ColumnRef, DataRow, Expr
-from pixeltable.io.external_store import Project, SyncStatus
+from pixeltable.io.external_store import Project
 from pixeltable.utils import coco
+from pixeltable.utils.local_store import TempStore
 
 # label_studio_sdk>=1 and label_studio_sdk<1 are not compatible, so we need to try
 # the import two different ways to insure intercompatibility
@@ -25,7 +28,7 @@ try:
     import label_studio_sdk.project as ls_project  # type: ignore
 except ImportError:
     # label_studio_sdk>=1 compatibility
-    import label_studio_sdk._legacy.project as ls_project  # type: ignore
+    import label_studio_sdk._legacy.project as ls_project
 
 _logger = logging.getLogger('pixeltable')
 
@@ -43,23 +46,26 @@ class LabelStudioProject(Project):
     """
     An [`ExternalStore`][pixeltable.io.ExternalStore] that represents a Label Studio project, providing functionality
     for synchronizing between a Pixeltable table and a Label Studio project.
+
+    The constructor will NOT create a new Label Studio project; it is also used when loading
+    metadata for existing projects.
     """
 
+    project_id: int  # Label Studio project ID
+    media_import_method: Literal['post', 'file', 'url']
+    _project: ls_project.Project | None
+
     def __init__(
         self,
         name: str,
         project_id: int,
         media_import_method: Literal['post', 'file', 'url'],
-        col_mapping: dict[Column, str],
-        stored_proxies: Optional[dict[Column, Column]] = None,
+        col_mapping: dict[ColumnHandle, str],
+        stored_proxies: dict[ColumnHandle, ColumnHandle] | None = None,
     ):
-        """
-        The constructor will NOT create a new Label Studio project; it is also used when loading
-        metadata for existing projects.
-        """
         self.project_id = project_id
         self.media_import_method = media_import_method
-        self._project: Optional[ls_project.Project] = None
+        self._project = None
         super().__init__(name, col_mapping, stored_proxies)
 
     @property
@@ -105,20 +111,20 @@ class LabelStudioProject(Project):
         """
         return {ANNOTATIONS_COLUMN: ts.JsonType(nullable=True)}
 
-    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> UpdateStatus:
         _logger.info(
             f'Syncing Label Studio project "{self.project_title}" with table `{t._name}`'
             f' (export: {export_data}, import: {import_data}).'
         )
         # Collect all existing tasks into a dict with entries `rowid: task`
         tasks = {tuple(task['meta']['rowid']): task for task in self.__fetch_all_tasks()}
-        sync_status = SyncStatus.empty()
+        sync_status = UpdateStatus()
         if export_data:
             export_sync_status = self.__update_tasks(t, tasks)
-            sync_status = sync_status.combine(export_sync_status)
+            sync_status += export_sync_status
         if import_data:
             import_sync_status = self.__update_table_from_tasks(t, tasks)
-            sync_status = sync_status.combine(import_sync_status)
+            sync_status += import_sync_status
         return sync_status
 
     def __fetch_all_tasks(self) -> Iterator[dict[str, Any]]:
@@ -142,7 +148,7 @@ class LabelStudioProject(Project):
                 f'Label Studio project {self.project_title!r}.'
             )
 
-    def __update_tasks(self, t: Table, existing_tasks: dict[tuple, dict]) -> SyncStatus:
+    def __update_tasks(self, t: Table, existing_tasks: dict[tuple, dict]) -> UpdateStatus:
         """
         Updates all tasks in this Label Studio project based on the Pixeltable data:
         - Creates new tasks for rows that don't map to any existing task;
@@ -155,7 +161,7 @@ class LabelStudioProject(Project):
         t_data_cols = [t_col for t_col, ext_col_name in self.col_mapping.items() if ext_col_name in config.data_keys]
 
         if len(t_data_cols) == 0:
-            return SyncStatus.empty()
+            return UpdateStatus()
 
         # Columns in `t` that map to `rectanglelabels` preannotations
         t_rl_cols = [
@@ -183,15 +189,15 @@ class LabelStudioProject(Project):
         self,
         t: Table,
         existing_tasks: dict[tuple, dict],
-        media_col: Column,
-        t_rl_cols: list[Column],
+        media_col: ColumnHandle,
+        t_rl_cols: list[ColumnHandle],
         rl_info: list['_RectangleLabel'],
-    ) -> SyncStatus:
-        is_stored = media_col.is_stored
+    ) -> UpdateStatus:
+        is_stored = media_col.get().is_stored
         # If it's a stored column, we can use `localpath`
-        localpath_col_opt = [t[media_col.name].localpath] if is_stored else []
+        localpath_col_opt = [t[media_col.get().name].localpath] if is_stored else []
         # Select the media column, rectanglelabels columns, and localpath (if appropriate)
-        rows = t.select(t[media_col.name], *[t[col.name] for col in t_rl_cols], *localpath_col_opt)
+        rows = t.select(t[media_col.get().name], *[t[col.get().name] for col in t_rl_cols], *localpath_col_opt)
         tasks_created = 0
         row_ids_in_pxt: set[tuple] = set()
 
@@ -209,7 +215,7 @@ class LabelStudioProject(Project):
                 else:
                     # No localpath; create a temp file and upload it
                     assert isinstance(row[media_col_idx], PIL.Image.Image)
-                    file = env.Env.get().create_tmp_path(extension='.png')
+                    file = TempStore.create_path(extension='.png')
                     row[media_col_idx].save(file, format='png')
                     task_id = self.project.import_tasks(file)[0]
                     os.remove(file)
@@ -232,48 +238,48 @@ class LabelStudioProject(Project):
 
         env.Env.get().console_logger.info(f'Created {tasks_created} new task(s) in {self}.')
 
-        sync_status = SyncStatus(external_rows_created=tasks_created)
+        sync_status = UpdateStatus(ext_row_count_stats=RowCountStats(ins_rows=tasks_created))
 
         deletion_sync_status = self.__delete_stale_tasks(existing_tasks, row_ids_in_pxt, tasks_created)
-
-        return sync_status.combine(deletion_sync_status)
+        sync_status += deletion_sync_status
+        return sync_status
 
     def __update_tasks_by_files(
         self,
         t: Table,
         existing_tasks: dict[tuple, dict],
-        t_data_cols: list[Column],
-        t_rl_cols: list[Column],
+        t_data_cols: list[ColumnHandle],
+        t_rl_cols: list[ColumnHandle],
         rl_info: list['_RectangleLabel'],
-    ) -> SyncStatus:
+    ) -> UpdateStatus:
         ext_data_cols = [self.col_mapping[col] for col in t_data_cols]
         expr_refs: dict[str, Expr] = {}  # kwargs for the select statement
         for col in t_data_cols:
-            col_name = col.name
+            col_name = col.get().name
             if self.media_import_method == 'url':
                 expr_refs[col_name] = t[col_name].fileurl
             else:
                 assert self.media_import_method == 'file'
-                if not col.col_type.is_media_type():
+                if not col.get().col_type.is_media_type():
                     # Not a media column; query the data directly
-                    expr_refs[col_name] = cast(ColumnRef, t[col_name])
+                    expr_refs[col_name] = t[col_name]
                 elif col in self.stored_proxies:
                     # Media column that has a stored proxy; use it. We have to give it a name,
                     # since it's an anonymous column
-                    stored_proxy_col = self.stored_proxies[col]
+                    stored_proxy_col = self.stored_proxies[col].get()
                     expr_refs[f'{col_name}_proxy'] = ColumnRef(stored_proxy_col).localpath
                 else:
                     # Media column without a stored proxy; this means it's a stored computed column,
                     # and we can just use the localpath
                     expr_refs[col_name] = t[col_name].localpath
 
-        df = t.select(*[t[col.name] for col in t_rl_cols], **expr_refs)
+        query = t.select(*[t[col.get().name] for col in t_rl_cols], **expr_refs)
         # The following buffers will hold `DataRow` indices that correspond to each of the selected
         # columns. `rl_col_idxs` holds the indices for the columns that map to RectangleLabels
         # preannotations; `data_col_idxs` holds the indices for the columns that map to data fields.
         # We have to wait until we begin iterating to populate them, so they're initially `None`.
-        rl_col_idxs: Optional[list[int]] = None
-        data_col_idxs: Optional[list[int]] = None
+        rl_col_idxs: list[int] | None = None
+        data_col_idxs: list[int] | None = None
 
         row_ids_in_pxt: set[tuple] = set()
         tasks_created = 0
@@ -286,11 +292,11 @@ class LabelStudioProject(Project):
             data_vals = [row[idx] for idx in data_col_idxs]
             coco_annotations = [row[idx] for idx in rl_col_idxs]
             for i in range(len(t_data_cols)):
-                if t_data_cols[i].col_type.is_media_type():
+                if t_data_cols[i].get().col_type.is_media_type():
                     # Special handling for media columns
                     assert isinstance(data_vals[i], str)
                     if self.media_import_method == 'url':
-                        data_vals[i] = self.__validate_fileurl(t_data_cols[i], data_vals[i])
+                        data_vals[i] = self.__validate_fileurl(t_data_cols[i].get(), data_vals[i])
                     else:
                         assert self.media_import_method == 'file'
                         data_vals[i] = self.__localpath_to_lspath(data_vals[i])
@@ -304,10 +310,10 @@ class LabelStudioProject(Project):
                 'predictions': predictions,
             }
 
-        for row in df._exec():
+        for row in query._exec():
             if rl_col_idxs is None:
-                rl_col_idxs = [expr.slot_idx for expr in df._select_list_exprs[: len(t_rl_cols)]]
-                data_col_idxs = [expr.slot_idx for expr in df._select_list_exprs[len(t_rl_cols) :]]
+                rl_col_idxs = [expr.slot_idx for expr in query._select_list_exprs[: len(t_rl_cols)]]
+                data_col_idxs = [expr.slot_idx for expr in query._select_list_exprs[len(t_rl_cols) :]]
             row_ids_in_pxt.add(row.rowid)
             task_info = create_task_info(row)
             # TODO(aaron-siegel): Implement more efficient update logic (currently involves a full table scan)
@@ -336,14 +342,14 @@ class LabelStudioProject(Project):
             f'Created {tasks_created} new task(s) and updated {tasks_updated} existing task(s) in {self}.'
         )
 
-        sync_status = SyncStatus(external_rows_created=tasks_created, external_rows_updated=tasks_updated)
+        sync_status = UpdateStatus(ext_row_count_stats=RowCountStats(ins_rows=tasks_created, upd_rows=tasks_updated))
 
         deletion_sync_status = self.__delete_stale_tasks(existing_tasks, row_ids_in_pxt, tasks_created)
-
-        return sync_status.combine(deletion_sync_status)
+        sync_status += deletion_sync_status
+        return sync_status
 
     @classmethod
-    def __validate_fileurl(cls, col: Column, url: str) -> Optional[str]:
+    def __validate_fileurl(cls, col: Column, url: str) -> str | None:
         # Check that the URL is one that will be visible to Label Studio. If it isn't, log an info message
         # to help users debug the issue.
         if not (url.startswith('http://') or url.startswith('https://')):
@@ -361,7 +367,7 @@ class LabelStudioProject(Project):
 
     def __delete_stale_tasks(
         self, existing_tasks: dict[tuple, dict], row_ids_in_pxt: set[tuple], tasks_created: int
-    ) -> SyncStatus:
+    ) -> UpdateStatus:
         deleted_rowids = set(existing_tasks.keys()) - row_ids_in_pxt
         # Sanity check the math
         assert len(deleted_rowids) == len(existing_tasks) + tasks_created - len(row_ids_in_pxt)
@@ -377,11 +383,11 @@ class LabelStudioProject(Project):
         for rowid in deleted_rowids:
             del existing_tasks[rowid]
 
-        return SyncStatus(external_rows_deleted=len(deleted_rowids))
+        return UpdateStatus(ext_row_count_stats=RowCountStats(del_rows=len(deleted_rowids)))
 
-    def __update_table_from_tasks(self, t: Table, tasks: dict[tuple, dict]) -> SyncStatus:
+    def __update_table_from_tasks(self, t: Table, tasks: dict[tuple, dict]) -> UpdateStatus:
         if ANNOTATIONS_COLUMN not in self.col_mapping.values():
-            return SyncStatus.empty()
+            return UpdateStatus()
 
         annotations = {
             # Replace [] by None to indicate no annotations. We do want to sync rows with no annotations,
@@ -391,7 +397,7 @@ class LabelStudioProject(Project):
             for task in tasks.values()
         }
 
-        local_annotations_col = next(k for k, v in self.col_mapping.items() if v == ANNOTATIONS_COLUMN)
+        local_annotations_col = next(k for k, v in self.col_mapping.items() if v == ANNOTATIONS_COLUMN).get()
 
         # Prune the annotations down to just the ones that have actually changed.
         rows = t.select(t[local_annotations_col.name])
@@ -412,23 +418,21 @@ class LabelStudioProject(Project):
             # TODO(aaron-siegel): Simplify this once propagation is properly implemented in batch_update
             ancestor = t
             while local_annotations_col not in ancestor._tbl_version.get().cols:
-                assert ancestor._base_table is not None
-                ancestor = ancestor._base_table
+                assert ancestor._get_base_table is not None
+                ancestor = ancestor._get_base_table()
             update_status = ancestor.batch_update(updates)
             env.Env.get().console_logger.info(f'Updated annotation(s) from {len(updates)} task(s) in {self}.')
-            return SyncStatus(pxt_rows_updated=update_status.num_rows, num_excs=update_status.num_excs)
+            return update_status
         else:
-            return SyncStatus.empty()
+            return UpdateStatus()
 
     def as_dict(self) -> dict[str, Any]:
         return {
             'name': self.name,
             'project_id': self.project_id,
             'media_import_method': self.media_import_method,
-            'col_mapping': [[self._column_as_dict(k), v] for k, v in self.col_mapping.items()],
-            'stored_proxies': [
-                [self._column_as_dict(k), self._column_as_dict(v)] for k, v in self.stored_proxies.items()
-            ],
+            'col_mapping': [[k.as_dict(), v] for k, v in self.col_mapping.items()],
+            'stored_proxies': [[k.as_dict(), v.as_dict()] for k, v in self.stored_proxies.items()],
         }
 
     @classmethod
@@ -437,8 +441,8 @@ class LabelStudioProject(Project):
             md['name'],
             md['project_id'],
             md['media_import_method'],
-            {cls._column_from_dict(entry[0]): entry[1] for entry in md['col_mapping']},
-            {cls._column_from_dict(entry[0]): cls._column_from_dict(entry[1]) for entry in md['stored_proxies']},
+            {ColumnHandle.from_dict(entry[0]): entry[1] for entry in md['col_mapping']},
+            {ColumnHandle.from_dict(entry[0]): ColumnHandle.from_dict(entry[1]) for entry in md['stored_proxies']},
         )
 
     def __repr__(self) -> str:
@@ -493,7 +497,7 @@ class LabelStudioProject(Project):
 
     @classmethod
     def __coco_to_predictions(
-        cls, coco_annotations: dict[str, Any], from_name: str, rl_info: '_RectangleLabel', task_id: Optional[int] = None
+        cls, coco_annotations: dict[str, Any], from_name: str, rl_info: '_RectangleLabel', task_id: int | None = None
     ) -> dict[str, Any]:
         width = coco_annotations['image']['width']
         height = coco_annotations['image']['height']
@@ -545,11 +549,11 @@ class LabelStudioProject(Project):
         cls,
         t: Table,
         label_config: str,
-        name: Optional[str],
-        title: Optional[str],
+        name: str | None,
+        title: str | None,
         media_import_method: Literal['post', 'file', 'url'],
-        col_mapping: Optional[dict[str, str]],
-        s3_configuration: Optional[dict[str, Any]],
+        col_mapping: dict[str, str] | None,
+        s3_configuration: dict[str, Any] | None,
         **kwargs: Any,
     ) -> 'LabelStudioProject':
         """
@@ -560,7 +564,7 @@ class LabelStudioProject(Project):
 
         if name is None:
             # Create a default name that's unique to the table
-            all_stores = t.external_stores
+            all_stores = t.external_stores()
             n = 0
             while f'ls_project_{n}' in all_stores:
                 n += 1
@@ -576,8 +580,8 @@ class LabelStudioProject(Project):
                 local_annotations_column = ANNOTATIONS_COLUMN
             else:
                 local_annotations_column = next(k for k, v in col_mapping.items() if v == ANNOTATIONS_COLUMN)
-            if local_annotations_column not in t._schema:
-                t.add_columns({local_annotations_column: ts.JsonType(nullable=True)})
+            if local_annotations_column not in t._get_schema():
+                t.add_columns({local_annotations_column: ts.Json})
 
         resolved_col_mapping = cls.validate_columns(
             t, config.export_columns, {ANNOTATIONS_COLUMN: ts.JsonType(nullable=True)}, col_mapping
@@ -648,7 +652,7 @@ class LabelStudioProject(Project):
 
 @dataclass(frozen=True)
 class _DataKey:
-    name: Optional[str]  # The 'name' attribute of the data key; may differ from the field name
+    name: str | None  # The 'name' attribute of the data key; may differ from the field name
     column_type: ts.ColumnType
 
 

pixeltable/io/lancedb.py

@@ -0,0 +1,3 @@
+from pixeltable.utils.lancedb import export_lancedb
+
+__all__ = ['export_lancedb']

pixeltable/io/pandas.py

@@ -1,5 +1,6 @@
 import os
-from typing import Any, Optional, Union
+import uuid
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -16,8 +17,8 @@ def import_pandas(
     tbl_name: str,
     df: pd.DataFrame,
     *,
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
 ) -> pxt.Table:
@@ -55,9 +56,9 @@ def import_pandas(
 
 def import_csv(
     tbl_name: str,
-    filepath_or_buffer: Union[str, os.PathLike],
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    filepath_or_buffer: str | os.PathLike,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
     **kwargs: Any,
@@ -84,10 +85,10 @@ def import_csv(
 
 def import_excel(
     tbl_name: str,
-    io: Union[str, os.PathLike],
+    io: str | os.PathLike,
     *,
-    schema_overrides: Optional[dict[str, Any]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
+    schema_overrides: dict[str, Any] | None = None,
+    primary_key: str | list[str] | None = None,
     num_retained_versions: int = 10,
     comment: str = '',
     **kwargs: Any,
@@ -132,6 +133,7 @@ def df_infer_schema(
     pd_schema: dict[str, ts.ColumnType] = {}
     for pd_name, pd_dtype in zip(df.columns, df.dtypes):
         if pd_name in schema_overrides:
+            assert isinstance(schema_overrides[pd_name], ts.ColumnType)
             pxt_type = schema_overrides[pd_name]
         else:
             pxt_type = __pd_coltype_to_pxt_type(pd_dtype, df[pd_name], pd_name not in primary_key)
@@ -140,7 +142,7 @@ def df_infer_schema(
     return pd_schema
 
 
-def __pd_dtype_to_pxt_type(pd_dtype: DtypeObj, nullable: bool) -> Optional[ts.ColumnType]:
+def __pd_dtype_to_pxt_type(pd_dtype: DtypeObj, nullable: bool) -> ts.ColumnType | None:
     """
     Determines a pixeltable ColumnType from a pandas dtype
 
@@ -159,7 +161,7 @@ def __pd_dtype_to_pxt_type(pd_dtype: DtypeObj, nullable: bool) -> Optional[ts.Co
         return None
     # Most other pandas dtypes are directly NumPy compatible
     assert isinstance(pd_dtype, np.dtype)
-    return ts.ArrayType.from_np_dtype(pd_dtype, nullable)
+    return ts.ColumnType.from_np_dtype(pd_dtype, nullable)
 
 
 def __pd_coltype_to_pxt_type(pd_dtype: DtypeObj, data_col: pd.Series, nullable: bool) -> ts.ColumnType:
@@ -191,7 +193,7 @@ def __pd_coltype_to_pxt_type(pd_dtype: DtypeObj, data_col: pd.Series, nullable:
 
 
 def _df_row_to_pxt_row(
-    row: tuple[Any, ...], schema: dict[str, ts.ColumnType], col_mapping: Optional[dict[str, str]]
+    row: tuple[Any, ...], schema: dict[str, ts.ColumnType], col_mapping: dict[str, str] | None
 ) -> dict[str, Any]:
     """Convert a row to insertable format"""
     pxt_row: dict[str, Any] = {}
@@ -229,6 +231,13 @@ def _df_row_to_pxt_row(
                     nval = pd.Timestamp(tval).tz_localize(tz=Env.get().default_time_zone)
                 else:
                     nval = tval.astimezone(Env.get().default_time_zone)
+        elif pxt_type.is_uuid_type():
+            if pd.isnull(val):
+                nval = None
+            elif isinstance(val, uuid.UUID):
+                nval = val
+            else:
+                nval = uuid.UUID(val)
         else:
             nval = val
         pxt_row[pxt_name] = nval