pixeltable 0.4.0rc3__py3-none-any.whl → 0.4.20__py3-none-any.whl

pixeltable/io/external_store.py

@@ -3,14 +3,13 @@ from __future__ import annotations
 import abc
 import itertools
 import logging
-from dataclasses import dataclass
-from typing import Any, Optional
-from uuid import UUID
+from typing import Any
 
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 from pixeltable import Column, Table
-from pixeltable.catalog import TableVersion
+from pixeltable.catalog import ColumnHandle, TableVersion
+from pixeltable.catalog.update_status import UpdateStatus
 
 _logger = logging.getLogger('pixeltable')
 
@@ -22,6 +21,8 @@ class ExternalStore(abc.ABC):
     and stateful external stores.
     """
 
+    __name: str
+
     def __init__(self, name: str) -> None:
         self.__name = name
 
@@ -38,13 +39,13 @@ class ExternalStore(abc.ABC):
         """Removes store-specific metadata created in link()."""
 
     @abc.abstractmethod
-    def get_local_columns(self) -> list[Column]:
+    def get_local_columns(self) -> list[ColumnHandle]:
         """
         Gets a list of all local (Pixeltable) columns that are associated with this external store.
         """
 
     @abc.abstractmethod
-    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> UpdateStatus:
         """
         Called by `Table.sync()` to implement store-specific synchronization logic.
         """
@@ -63,9 +64,12 @@ class Project(ExternalStore, abc.ABC):
     additional capabilities specific to such projects.
     """
 
-    stored_proxies: dict[Column, Column]
+    _col_mapping: dict[ColumnHandle, str]  # col -> external col name
+    stored_proxies: dict[ColumnHandle, ColumnHandle]  # original col -> proxy col
 
-    def __init__(self, name: str, col_mapping: dict[Column, str], stored_proxies: Optional[dict[Column, Column]]):
+    def __init__(
+        self, name: str, col_mapping: dict[ColumnHandle, str], stored_proxies: dict[ColumnHandle, ColumnHandle] | None
+    ):
         super().__init__(name)
         self._col_mapping = col_mapping
 
@@ -80,11 +84,11 @@ class Project(ExternalStore, abc.ABC):
         # Note from aaron-siegel: This methodology is inefficient in the case where a table has many views with a high
         # proportion of overlapping rows, all proxying the same base column.
         if stored_proxies is None:
-            self.stored_proxies: dict[Column, Column] = {}
+            self.stored_proxies: dict[ColumnHandle, ColumnHandle] = {}
         else:
             self.stored_proxies = stored_proxies
 
-    def get_local_columns(self) -> list[Column]:
+    def get_local_columns(self) -> list[ColumnHandle]:
         return list(self.col_mapping.keys())
 
     def link(self, tbl_version: TableVersion) -> None:
@@ -92,15 +96,16 @@ class Project(ExternalStore, abc.ABC):
         # This ensures that the media in those columns resides in the media store.
         # First determine which columns (if any) need stored proxies, but don't have one yet.
         stored_proxies_needed: list[Column] = []
-        for col in self.col_mapping:
+        for col_handle in self.col_mapping:
+            col = col_handle.get()
             if col.col_type.is_media_type() and not (col.is_stored and col.is_computed):
                 # If this column is already proxied in some other Project, use the existing proxy to avoid
                 # duplication. Otherwise, we'll create a new one.
                 for store in tbl_version.external_stores.values():
-                    if isinstance(store, Project) and col in store.stored_proxies:
-                        self.stored_proxies[col] = store.stored_proxies[col]
+                    if isinstance(store, Project) and col_handle in store.stored_proxies:
+                        self.stored_proxies[col_handle] = store.stored_proxies[col_handle]
                         break
-                if col not in self.stored_proxies:
+                if col_handle not in self.stored_proxies:
                     # We didn't find it in an existing Project
                     stored_proxies_needed.append(col)
 
@@ -110,17 +115,20 @@ class Project(ExternalStore, abc.ABC):
             proxy_cols = [self.create_stored_proxy(col) for col in stored_proxies_needed]
             # Add the columns; this will also update table metadata.
             tbl_version.add_columns(proxy_cols, print_stats=False, on_error='ignore')
+            self.stored_proxies.update(
+                {col.handle: proxy_col.handle for col, proxy_col in zip(stored_proxies_needed, proxy_cols)}
+            )
 
     def unlink(self, tbl_version: TableVersion) -> None:
         # Determine which stored proxies can be deleted. (A stored proxy can be deleted if it is not referenced by
         # any *other* external store for this table.)
-        deletions_needed: set[Column] = set(self.stored_proxies.values())
+        deletions_needed: set[ColumnHandle] = set(self.stored_proxies.values())
         for name, store in tbl_version.external_stores.items():
             if isinstance(store, Project) and name != self.name:
                 deletions_needed = deletions_needed.difference(set(store.stored_proxies.values()))
         if len(deletions_needed) > 0:
-            _logger.info(f'Removing stored proxies for columns: {[col.name for col in deletions_needed]}')
-            tbl_version._drop_columns(deletions_needed)
+            _logger.info(f'Removing stored proxies for columns: {[col.get().name for col in deletions_needed]}')
+            tbl_version._drop_columns(col.get() for col in deletions_needed)
             self.stored_proxies.clear()
 
     def create_stored_proxy(self, col: Column) -> Column:
@@ -142,11 +150,10 @@ class Project(ExternalStore, abc.ABC):
             computed_with=exprs.ColumnRef(col).apply(lambda x: x, col_type=col.col_type),
             stored=True,
         )
-        self.stored_proxies[col] = proxy_col
         return proxy_col
 
     @property
-    def col_mapping(self) -> dict[Column, str]:
+    def col_mapping(self) -> dict[ColumnHandle, str]:
         return self._col_mapping
 
     @abc.abstractmethod
@@ -180,8 +187,8 @@ class Project(ExternalStore, abc.ABC):
         table: Table,
         export_cols: dict[str, ts.ColumnType],
         import_cols: dict[str, ts.ColumnType],
-        col_mapping: Optional[dict[str, str]],
-    ) -> dict[Column, str]:
+        col_mapping: dict[str, str] | None,
+    ) -> dict[ColumnHandle, str]:
         """
         Verifies that the specified `col_mapping` is valid. In particular, checks that:
         (i) the keys of `col_mapping` are valid columns of the specified `Table`;
@@ -199,33 +206,34 @@ class Project(ExternalStore, abc.ABC):
         if col_mapping is None:
             col_mapping = {col: col for col in itertools.chain(export_cols.keys(), import_cols.keys())}
 
-        resolved_col_mapping: dict[Column, str] = {}
+        resolved_col_mapping: dict[ColumnHandle, str] = {}
 
         # Validate names
-        t_cols = set(table._schema.keys())
+        t_cols = set(table._get_schema().keys())
         for t_col, ext_col in col_mapping.items():
             if t_col not in t_cols:
                 if is_user_specified_col_mapping:
                     raise excs.Error(
-                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table._name}` '
+                        f'Column name {t_col!r} appears as a key in `col_mapping`, but {table._display_str()} '
                         'contains no such column.'
                     )
                 else:
                     raise excs.Error(
-                        f'Column `{t_col}` does not exist in Table `{table._name}`. Either add a column `{t_col}`, '
+                        f'Column {t_col!r} does not exist in {table._display_str()}. Either add a column {t_col!r}, '
                         f'or specify a `col_mapping` to associate a different column with '
-                        f'the external field `{ext_col}`.'
+                        f'the external field {ext_col!r}.'
                     )
             if ext_col not in export_cols and ext_col not in import_cols:
                 raise excs.Error(
-                    f'Column name `{ext_col}` appears as a value in `col_mapping`, but the external store '
-                    f'configuration has no column `{ext_col}`.'
+                    f'Column name {ext_col!r} appears as a value in `col_mapping`, but the external store '
+                    f'configuration has no column {ext_col!r}.'
                 )
             col_ref = table[t_col]
             assert isinstance(col_ref, exprs.ColumnRef)
-            resolved_col_mapping[col_ref.col] = ext_col
+            resolved_col_mapping[col_ref.col.handle] = ext_col
+
         # Validate column specs
-        t_col_types = table._schema
+        t_col_types = table._get_schema()
         for t_col, ext_col in col_mapping.items():
             t_col_type = t_col_types[t_col]
             if ext_col in export_cols:
@@ -233,57 +241,23 @@ class Project(ExternalStore, abc.ABC):
                 ext_col_type = export_cols[ext_col]
                 if not ext_col_type.is_supertype_of(t_col_type, ignore_nullable=True):
                     raise excs.Error(
-                        f'Column `{t_col}` cannot be exported to external column `{ext_col}` '
+                        f'Column {t_col!r} cannot be exported to external column {ext_col!r} '
                         f'(incompatible types; expecting `{ext_col_type}`)'
                     )
             if ext_col in import_cols:
                 # Validate that the external column can be assigned to the table column
                 if table._tbl_version_path.get_column(t_col).is_computed:
                     raise excs.Error(
-                        f'Column `{t_col}` is a computed column, which cannot be populated from an external column'
+                        f'Column {t_col!r} is a computed column, which cannot be populated from an external column'
                     )
                 ext_col_type = import_cols[ext_col]
                 if not t_col_type.is_supertype_of(ext_col_type, ignore_nullable=True):
                     raise excs.Error(
-                        f'Column `{t_col}` cannot be imported from external column `{ext_col}` '
+                        f'Column {t_col!r} cannot be imported from external column {ext_col!r} '
                         f'(incompatible types; expecting `{ext_col_type}`)'
                     )
         return resolved_col_mapping
 
-    @classmethod
-    def _column_as_dict(cls, col: Column) -> dict[str, Any]:
-        return {'tbl_id': str(col.tbl.id), 'col_id': col.id}
-
-    @classmethod
-    def _column_from_dict(cls, d: dict[str, Any]) -> Column:
-        from pixeltable.catalog import Catalog
-
-        tbl_id = UUID(d['tbl_id'])
-        col_id = d['col_id']
-        return Catalog.get().get_tbl_version(tbl_id, None).cols_by_id[col_id]
-
-
-@dataclass(frozen=True)
-class SyncStatus:
-    external_rows_created: int = 0
-    external_rows_deleted: int = 0
-    external_rows_updated: int = 0
-    pxt_rows_updated: int = 0
-    num_excs: int = 0
-
-    def combine(self, other: 'SyncStatus') -> 'SyncStatus':
-        return SyncStatus(
-            external_rows_created=self.external_rows_created + other.external_rows_created,
-            external_rows_deleted=self.external_rows_deleted + other.external_rows_deleted,
-            external_rows_updated=self.external_rows_updated + other.external_rows_updated,
-            pxt_rows_updated=self.pxt_rows_updated + other.pxt_rows_updated,
-            num_excs=self.num_excs + other.num_excs,
-        )
-
-    @classmethod
-    def empty(cls) -> 'SyncStatus':
-        return SyncStatus(0, 0, 0, 0, 0)
-
 
 class MockProject(Project):
     """A project that cannot be synced, used mainly for testing."""
@@ -293,8 +267,8 @@ class MockProject(Project):
         name: str,
         export_cols: dict[str, ts.ColumnType],
         import_cols: dict[str, ts.ColumnType],
-        col_mapping: dict[Column, str],
-        stored_proxies: Optional[dict[Column, Column]] = None,
+        col_mapping: dict[ColumnHandle, str],
+        stored_proxies: dict[ColumnHandle, ColumnHandle] | None = None,
     ):
         super().__init__(name, col_mapping, stored_proxies)
         self.export_cols = export_cols
@@ -308,7 +282,7 @@ class MockProject(Project):
         name: str,
         export_cols: dict[str, ts.ColumnType],
         import_cols: dict[str, ts.ColumnType],
-        col_mapping: Optional[dict[str, str]] = None,
+        col_mapping: dict[str, str] | None = None,
     ) -> 'MockProject':
         col_mapping = cls.validate_columns(t, export_cols, import_cols, col_mapping)
         return cls(name, export_cols, import_cols, col_mapping)
@@ -319,7 +293,7 @@ class MockProject(Project):
     def get_import_columns(self) -> dict[str, ts.ColumnType]:
         return self.import_cols
 
-    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> UpdateStatus:
         raise NotImplementedError()
 
     def delete(self) -> None:
@@ -334,10 +308,8 @@ class MockProject(Project):
             'name': self.name,
             'export_cols': {k: v.as_dict() for k, v in self.export_cols.items()},
             'import_cols': {k: v.as_dict() for k, v in self.import_cols.items()},
-            '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
@@ -346,8 +318,8 @@ class MockProject(Project):
             md['name'],
             {k: ts.ColumnType.from_dict(v) for k, v in md['export_cols'].items()},
             {k: ts.ColumnType.from_dict(v) for k, v in md['import_cols'].items()},
-            {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 __eq__(self, other: object) -> bool:

pixeltable/io/fiftyone.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Iterator, Optional, Union
+from typing import Any, Iterator
 
 import fiftyone as fo  # type: ignore[import-untyped]
 import fiftyone.utils.data as foud  # type: ignore[import-untyped]
@@ -9,7 +9,7 @@ import puremagic
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 from pixeltable import exprs
-from pixeltable.env import Env
+from pixeltable.utils.local_store import TempStore
 
 
 class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
@@ -20,7 +20,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
     __image_format: str  # format to use for any exported images that are not already stored on disk
     __labels: dict[str, tuple[exprs.Expr, type[fo.Label]]]  # label_name -> (expr, label_cls)
     __image_idx: int  # index of the image expr in the select list
-    __localpath_idx: Optional[int]  # index of the image localpath in the select list, if present
+    __localpath_idx: int | None  # index of the image localpath in the select list, if present
     __row_iter: Iterator[list]  # iterator over the table rows, to be convered to FiftyOne samples
 
     def __init__(
@@ -28,12 +28,12 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         tbl: pxt.Table,
         image: exprs.Expr,
         image_format: str,
-        classifications: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
-        detections: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
-        dataset_dir: Optional[os.PathLike] = None,
+        classifications: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
+        detections: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
+        dataset_dir: os.PathLike | None = None,
         shuffle: bool = False,
-        seed: Union[int, float, str, bytes, bytearray, None] = None,
-        max_samples: Optional[int] = None,
+        seed: int | float | str | bytes | bytearray | None = None,
+        max_samples: int | None = None,
     ):
         super().__init__(dataset_dir=dataset_dir, shuffle=shuffle, seed=seed, max_samples=max_samples)
 
@@ -90,7 +90,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
         df = tbl.select(*selection)
         self.__row_iter = df._output_row_iterator()
 
-    def __next__(self) -> tuple[str, Optional[fo.ImageMetadata], Optional[dict[str, fo.Label]]]:
+    def __next__(self) -> tuple[str, fo.ImageMetadata | None, dict[str, fo.Label] | None]:
         row = next(self.__row_iter)
         img = row[self.__image_idx]
         assert isinstance(img, PIL.Image.Image)
@@ -100,7 +100,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             assert isinstance(file, str)
         else:
             # Write the dynamically created image to a temp file
-            file = str(Env.get().create_tmp_path(f'.{self.__image_format}'))
+            file = TempStore.create_path(extension=f'.{self.__image_format}')
             img.save(file, format=self.__image_format)
 
         metadata = fo.ImageMetadata(
@@ -108,7 +108,7 @@ class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
             mime_type=puremagic.from_file(file, mime=True),
             width=img.width,
             height=img.height,
-            filepath=file,
+            filepath=str(file),
             num_channels=len(img.getbands()),
         )
 

pixeltable/io/globals.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
 from pixeltable import Table, exprs
+from pixeltable.catalog.update_status import UpdateStatus
 from pixeltable.env import Env
-from pixeltable.io.external_store import SyncStatus
 
 if TYPE_CHECKING:
     import fiftyone as fo  # type: ignore[import-untyped]
@@ -15,14 +15,14 @@ if TYPE_CHECKING:
 def create_label_studio_project(
     t: Table,
     label_config: str,
-    name: Optional[str] = None,
-    title: Optional[str] = None,
+    name: str | None = None,
+    title: str | None = None,
     media_import_method: Literal['post', 'file', 'url'] = 'post',
-    col_mapping: Optional[dict[str, str]] = None,
+    col_mapping: dict[str, str] | None = None,
     sync_immediately: bool = True,
-    s3_configuration: Optional[dict[str, Any]] = None,
+    s3_configuration: dict[str, Any] | None = None,
     **kwargs: Any,
-) -> SyncStatus:
+) -> UpdateStatus:
     """
     Create a new Label Studio project and link it to the specified [`Table`][pixeltable.Table].
 
@@ -96,32 +96,33 @@ def create_label_studio_project(
             [Label Studio start_project docs](https://labelstud.io/sdk/project.html#label_studio_sdk.project.Project.start_project).
 
     Returns:
-        A `SyncStatus` representing the status of any synchronization operations that occurred.
+        An `UpdateStatus` representing the status of any synchronization operations that occurred.
 
     Examples:
         Create a Label Studio project whose tasks correspond to videos stored in the `video_col`
         column of the table `tbl`:
 
         >>> config = \"\"\"
-            <View>
-                <Video name="video_obj" value="$video_col"/>
-                <Choices name="video-category" toName="video" showInLine="true">
-                    <Choice value="city"/>
-                    <Choice value="food"/>
-                    <Choice value="sports"/>
-                </Choices>
-            </View>\"\"\"
-            create_label_studio_project(tbl, config)
+        ... <View>
+        ...     <Video name="video_obj" value="$video_col"/>
+        ...     <Choices name="video-category" toName="video" showInLine="true">
+        ...         <Choice value="city"/>
+        ...         <Choice value="food"/>
+        ...         <Choice value="sports"/>
+        ...     </Choices>
+        ... </View>
+        ... \"\"\"
+        >>> create_label_studio_project(tbl, config)
 
         Create a Label Studio project with the same configuration, using `media_import_method='url'`,
         whose media are stored in an S3 bucket:
 
         >>> create_label_studio_project(
-                tbl,
-                config,
-                media_import_method='url',
-                s3_configuration={'bucket': 'my-bucket', 'region_name': 'us-east-2'}
-            )
+        ...     tbl,
+        ...     config,
+        ...     media_import_method='url',
+        ...     s3_configuration={'bucket': 'my-bucket', 'region_name': 'us-east-2'}
+        ... )
     """
     Env.get().require_package('label_studio_sdk')
 
@@ -136,22 +137,22 @@ def create_label_studio_project(
     if sync_immediately:
         return t.sync()
     else:
-        return SyncStatus.empty()
+        return UpdateStatus()
 
 
 def export_images_as_fo_dataset(
     tbl: pxt.Table,
     images: exprs.Expr,
     image_format: str = 'webp',
-    classifications: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
-    detections: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
+    classifications: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
+    detections: exprs.Expr | list[exprs.Expr] | dict[str, exprs.Expr] | None = None,
 ) -> 'fo.Dataset':
     """
     Export images from a Pixeltable table as a Voxel51 dataset. The data must consist of a single column
     (or expression) containing image data, along with optional additional columns containing labels. Currently, only
     classification and detection labels are supported.
 
-    The [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/docs/working-with-voxel51) tutorial contains a
+    The [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/examples/vision/voxel51) tutorial contains a
     fully worked example showing how to export data from a Pixeltable table and load it into Voxel51.
 
     Images in the dataset that already exist on disk will be exported directly, in whatever format they
@@ -204,13 +205,13 @@ def export_images_as_fo_dataset(
         Export the images in the `image` column of the table `tbl` as a Voxel51 dataset, using classification
         labels from `tbl.classifications`:
 
-        >>> export_as_fiftyone(
+        >>> export_images_as_fo_dataset(
         ...     tbl,
         ...     tbl.image,
         ...     classifications=tbl.classifications
         ... )
 
-        See the [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/docs/working-with-voxel51) tutorial
+        See the [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/examples/vision/voxel51) tutorial
         for a fully worked example.
     """
     Env.get().require_package('fiftyone')

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
@@ -36,7 +36,7 @@ _hf_to_pxt: dict[str, ts.ColumnType] = {
 }
 
 
-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
 
@@ -50,15 +50,23 @@ def _to_pixeltable_type(feature_type: Any, nullable: bool) -> Optional[ts.Column
     elif isinstance(feature_type, datasets.Sequence):
         # 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, 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
 
@@ -68,7 +76,7 @@ def _get_hf_schema(dataset: Union[datasets.Dataset, datasets.DatasetDict]) -> da
 
 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,10 +91,10 @@ def huggingface_schema_to_pxt_schema(
 
 def import_huggingface_dataset(
     table_path: str,
-    dataset: Union[datasets.Dataset, datasets.DatasetDict],
+    dataset: datasets.Dataset | datasets.DatasetDict,
     *,
-    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.