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

pixeltable/io/globals.py

@@ -1,31 +1,33 @@
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+from __future__ import annotations
+
+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]
 
 
 def create_label_studio_project(
-        t: Table,
-        label_config: str,
-        name: Optional[str] = None,
-        title: Optional[str] = None,
-        media_import_method: Literal['post', 'file', 'url'] = 'post',
-        col_mapping: Optional[dict[str, str]] = None,
-        sync_immediately: bool = True,
-        s3_configuration: Optional[dict[str, Any]] = None,
-        **kwargs: Any
-) -> SyncStatus:
+    t: Table,
+    label_config: str,
+    name: str | None = None,
+    title: str | None = None,
+    media_import_method: Literal['post', 'file', 'url'] = 'post',
+    col_mapping: dict[str, str] | None = None,
+    sync_immediately: bool = True,
+    s3_configuration: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> UpdateStatus:
     """
     Create a new Label Studio project and link it to the specified [`Table`][pixeltable.Table].
 
     - A tutorial notebook with fully worked examples can be found here:
-      [Using Label Studio for Annotations with Pixeltable](https://pixeltable.readme.io/docs/label-studio)
+      [Using Label Studio for Annotations with Pixeltable](https://docs.pixeltable.com/notebooks/integrations/using-label-studio-with-pixeltable)
 
     The required parameter `label_config` specifies the Label Studio project configuration,
     in XML format, as described in the Label Studio documentation. The linked project will
@@ -86,53 +88,48 @@ def create_label_studio_project(
             parameters of the Label Studio `connect_s3_import_storage` method, as described in the
             [Label Studio connect_s3_import_storage docs](https://labelstud.io/sdk/project.html#label_studio_sdk.project.Project.connect_s3_import_storage).
             `bucket` must be specified; all other parameters are optional. If credentials are not specified explicitly,
-            Pixeltable will attempt to retrieve them from the environment (such as from `~/.aws/credentials`). If a title is not
-            specified, Pixeltable will use the default `'Pixeltable-S3-Import-Storage'`. All other parameters use their Label
-            Studio defaults.
+            Pixeltable will attempt to retrieve them from the environment (such as from `~/.aws/credentials`).
+            If a title is not specified, Pixeltable will use the default `'Pixeltable-S3-Import-Storage'`.
+            All other parameters use their Label Studio defaults.
         kwargs: Additional keyword arguments are passed to the `start_project` method in the Label
             Studio SDK, as described in the
             [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`:
+        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')
 
     from pixeltable.io.label_studio import LabelStudioProject
 
     ls_project = LabelStudioProject.create(
-        t,
-        label_config,
-        name,
-        title,
-        media_import_method,
-        col_mapping,
-        s3_configuration,
-        **kwargs
+        t, label_config, name, title, media_import_method, col_mapping, s3_configuration, **kwargs
     )
 
     # Link the project to `t`, and sync if appropriate.
@@ -140,159 +137,27 @@ def create_label_studio_project(
     if sync_immediately:
         return t.sync()
     else:
-        return SyncStatus.empty()
-
-
-def import_rows(
-    tbl_path: str,
-    rows: list[dict[str, Any]],
-    *,
-    schema_overrides: Optional[dict[str, pxt.ColumnType]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
-    num_retained_versions: int = 10,
-    comment: str = ''
-    ) -> Table:
-    """
-    Creates a new base table from a list of dictionaries. The dictionaries must be of the
-    form `{column_name: value, ...}`. Pixeltable will attempt to infer the schema of the table from the
-    supplied data, using the most specific type that can represent all the values in a column.
-
-    If `schema_overrides` is specified, then for each entry `(column_name, type)` in `schema_overrides`,
-    Pixeltable will force the specified column to the specified type (and will not attempt any type inference
-    for that column).
-
-    All column types of the new table will be nullable unless explicitly specified as non-nullable in
-    `schema_overrides`.
-
-    Args:
-        tbl_path: The qualified name of the table to create.
-        rows: The list of dictionaries to import.
-        schema_overrides: If specified, then columns in `schema_overrides` will be given the specified types
-            as described above.
-        primary_key: The primary key of the table (see [`create_table()`][pixeltable.create_table]).
-        num_retained_versions: The number of retained versions of the table (see [`create_table()`][pixeltable.create_table]).
-        comment: A comment to attach to the table (see [`create_table()`][pixeltable.create_table]).
-
-    Returns:
-        A handle to the newly created [`Table`][pixeltable.Table].
-    """
-    if schema_overrides is None:
-        schema_overrides = {}
-    schema: dict[str, pxt.ColumnType] = {}
-    cols_with_nones: set[str] = set()
-
-    for n, row in enumerate(rows):
-        for col_name, value in row.items():
-            if col_name in schema_overrides:
-                # We do the insertion here; this will ensure that the column order matches the order
-                # in which the column names are encountered in the input data, even if `schema_overrides`
-                # is specified.
-                if col_name not in schema:
-                    schema[col_name] = schema_overrides[col_name]
-            elif value is not None:
-                # If `key` is not in `schema_overrides`, then we infer its type from the data.
-                # The column type will always be nullable by default.
-                col_type = pxt.ColumnType.infer_literal_type(value, nullable=True)
-                if col_type is None:
-                    raise excs.Error(f'Could not infer type for column `{col_name}`; the value in row {n} has an unsupported type: {type(value)}')
-                if col_name not in schema:
-                    schema[col_name] = col_type
-                else:
-                    supertype = schema[col_name].supertype(col_type)
-                    if supertype is None:
-                        raise excs.Error(
-                            f'Could not infer type of column `{col_name}`; the value in row {n} does not match preceding type {schema[col_name]}: {value!r}\n'
-                            'Consider specifying the type explicitly in `schema_overrides`.'
-                        )
-                    schema[col_name] = supertype
-            else:
-                cols_with_nones.add(col_name)
-
-    extraneous_keys = schema_overrides.keys() - schema.keys()
-    if len(extraneous_keys) > 0:
-        raise excs.Error(f'The following columns specified in `schema_overrides` are not present in the data: {", ".join(extraneous_keys)}')
-
-    entirely_none_cols = cols_with_nones - schema.keys()
-    if len(entirely_none_cols) > 0:
-        # A column can only end up in `entirely_null_cols` if it was not in `schema_overrides` and
-        # was not encountered in any row with a non-None value.
-        raise excs.Error(
-            f'The following columns have no non-null values: {", ".join(entirely_none_cols)}\n'
-            'Consider specifying the type(s) explicitly in `schema_overrides`.'
-        )
-
-    t = pxt.create_table(tbl_path, schema, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment)
-    t.insert(rows)
-    return t
-
-
-def import_json(
-    tbl_path: str,
-    filepath_or_url: str,
-    *,
-    schema_overrides: Optional[dict[str, pxt.ColumnType]] = None,
-    primary_key: Optional[Union[str, list[str]]] = None,
-    num_retained_versions: int = 10,
-    comment: str = '',
-    **kwargs: Any
-) -> Table:
-    """
-    Creates a new base table from a JSON file. This is a convenience method and is
-    equivalent to calling `import_data(table_path, json.loads(file_contents, **kwargs), ...)`, where `file_contents`
-    is the contents of the specified `filepath_or_url`.
-
-    Args:
-        tbl_path: The name of the table to create.
-        filepath_or_url: The path or URL of the JSON file.
-        schema_overrides: If specified, then columns in `schema_overrides` will be given the specified types
-            (see [`import_rows()`][pixeltable.io.import_rows]).
-        primary_key: The primary key of the table (see [`create_table()`][pixeltable.create_table]).
-        num_retained_versions: The number of retained versions of the table (see [`create_table()`][pixeltable.create_table]).
-        comment: A comment to attach to the table (see [`create_table()`][pixeltable.create_table]).
-        kwargs: Additional keyword arguments to pass to `json.loads`.
-
-    Returns:
-        A handle to the newly created [`Table`][pixeltable.Table].
-    """
-    import json
-    import urllib.parse
-    import urllib.request
-
-    # TODO Consolidate this logic with other places where files/URLs are parsed
-    parsed = urllib.parse.urlparse(filepath_or_url)
-    if len(parsed.scheme) <= 1 or parsed.scheme == 'file':
-        # local file path
-        if len(parsed.scheme) <= 1:
-            filepath = filepath_or_url
-        else:
-            filepath = urllib.parse.unquote(urllib.request.url2pathname(parsed.path))
-        with open(filepath) as fp:
-            contents = fp.read()
-    else:
-        # URL
-        contents = urllib.request.urlopen(filepath_or_url).read()
-    data = json.loads(contents, **kwargs)
-    return import_rows(tbl_path, data, schema_overrides=schema_overrides, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment)
+        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
     are stored in. Images that are not already on disk (such as frames extracted using a
-    [`FrameIterator`][pixeltable.iterators.FrameIterator]) will first be written to disk in the specified
+    [`frame_iterator`][pixeltable.functions.video.frame_iterator]) will first be written to disk in the specified
     `image_format`.
 
     The label parameters accept one or more sets of labels of each type. If a single `Expr` is provided, then it will
@@ -340,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')
@@ -358,6 +223,6 @@ def export_images_as_fo_dataset(
     if not images.col_type.is_image_type():
         raise excs.Error(f'`images` must be an expression of type Image (got {images.col_type._to_base_str()})')
 
-    return fo.Dataset.from_importer(PxtImageDatasetImporter(
-        tbl, images, image_format, classifications=classifications, detections=detections
-    ))
+    return fo.Dataset.from_importer(
+        PxtImageDatasetImporter(tbl, images, image_format, classifications=classifications, detections=detections)
+    )

pixeltable/io/hf_datasets.py

@@ -1,190 +1,159 @@
 from __future__ import annotations
 
-import logging
-import math
-import random
 import typing
-from typing import Union, Optional, Any
+from typing import Any
 
 import pixeltable as pxt
 import pixeltable.type_system as ts
-from pixeltable import exceptions as excs
 
 if typing.TYPE_CHECKING:
     import datasets  # type: ignore[import-untyped]
 
-_logger = logging.getLogger(__name__)
 
-# use 100MB as the batch size limit for loading a huggingface dataset into pixeltable.
-# The primary goal is to bound memory use, regardless of dataset size.
-# Second goal is to limit overhead. 100MB is presumed to be reasonable for a lot of storage systems.
-_K_BATCH_SIZE_BYTES = 100_000_000
-
-# note, there are many more types. we allow overrides in the schema_override parameter
+# note, there are many more types. we allow overrides in the schema_overrides parameter
 # to handle cases where the appropriate type is not yet mapped, or to override this mapping.
 # https://huggingface.co/docs/datasets/v2.17.0/en/package_reference/main_classes#datasets.Value
 _hf_to_pxt: dict[str, ts.ColumnType] = {
-    'int32': ts.IntType(nullable=True),  # pixeltable widens to big int
-    'int64': ts.IntType(nullable=True),
     'bool': ts.BoolType(nullable=True),
+    'int8': ts.IntType(nullable=True),
+    'int16': ts.IntType(nullable=True),
+    'int32': ts.IntType(nullable=True),
+    'int64': ts.IntType(nullable=True),
+    'uint8': ts.IntType(nullable=True),
+    'uint16': ts.IntType(nullable=True),
+    'uint32': ts.IntType(nullable=True),
+    'uint64': ts.IntType(nullable=True),
+    'float16': ts.FloatType(nullable=True),
     'float32': ts.FloatType(nullable=True),
+    'float64': ts.FloatType(nullable=True),
     'string': ts.StringType(nullable=True),
+    'large_string': ts.StringType(nullable=True),
     '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) -> 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
 
     if isinstance(feature_type, datasets.ClassLabel):
         # enum, example: ClassLabel(names=['neg', 'pos'], id=None)
-        return ts.StringType(nullable=True)
+        return ts.StringType(nullable=nullable)
     elif isinstance(feature_type, datasets.Value):
         # example: Value(dtype='int64', id=None)
-        return _hf_to_pxt.get(feature_type.dtype, None)
-    elif isinstance(feature_type, datasets.Sequence):
+        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, datasets.LargeList)):
         # example: cohere wiki. Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None)
-        dtype = _to_pixeltable_type(feature_type.feature)
-        length = feature_type.length if feature_type.length != -1 else None
-        return ts.ArrayType(shape=(length,), dtype=dtype)
+        dtype = _to_pixeltable_type(feature_type.feature, nullable)
+        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=True)
+        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_pixeltable_schema(
-    hf_dataset: Union[datasets.Dataset, datasets.DatasetDict],
-) -> dict[str, Optional[ts.ColumnType]]:
+def huggingface_schema_to_pxt_schema(
+    hf_schema: datasets.Features, schema_overrides: dict[str, Any], primary_key: list[str]
+) -> dict[str, ts.ColumnType | None]:
     """Generate a pixeltable schema from a huggingface dataset schema.
     Columns without a known mapping are mapped to None
     """
-    hf_schema = _get_hf_schema(hf_dataset)
     pixeltable_schema = {
-        column_name: _to_pixeltable_type(feature_type) for column_name, feature_type in hf_schema.items()
+        column_name: _to_pixeltable_type(feature_type, column_name not in primary_key)
+        if column_name not in schema_overrides
+        else schema_overrides[column_name]
+        for column_name, feature_type in hf_schema.items()
     }
     return pixeltable_schema
 
 
 def import_huggingface_dataset(
     table_path: str,
-    dataset: Union[datasets.Dataset, datasets.DatasetDict],
+    dataset: datasets.Dataset | datasets.DatasetDict | datasets.IterableDataset | datasets.IterableDatasetDict,
     *,
-    column_name_for_split: Optional[str] = None,
-    schema_overrides: Optional[dict[str, Any]] = 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.
-        column_name_for_split: column name to use for split information. If None, no split information will be stored.
+        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 they are valid
-            Pixeltable identifiers).
+            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
+            they are valid Pixeltable identifiers).
+        primary_key: The primary key of the table (see [`create_table()`][pixeltable.create_table]).
         kwargs: Additional arguments to pass to `create_table`.
+            An argument of `column_name_for_split` must be provided if the source is a DatasetDict.
+            This column name will contain the split information. If None, no split information will be stored.
 
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
     """
-    import datasets
-    import pixeltable as pxt
-
-    if table_path in pxt.list_tables():
-        raise excs.Error(f'table {table_path} already exists')
-
-    if not isinstance(dataset, (datasets.Dataset, datasets.DatasetDict)):
-        raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
-
-    if isinstance(dataset, datasets.Dataset):
-        # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
-        raw_name = dataset.split._name
-        split_name = raw_name.split('[')[0] if raw_name is not None else None
-        dataset_dict = {split_name: dataset}
-    else:
-        dataset_dict = dataset
-
-    pixeltable_schema = huggingface_schema_to_pixeltable_schema(dataset)
-    if schema_overrides is not None:
-        pixeltable_schema.update(schema_overrides)
-
-    if column_name_for_split is not None:
-        if column_name_for_split in pixeltable_schema:
-            raise excs.Error(
-                f'Column name `{column_name_for_split}` already exists in dataset schema; provide a different `column_name_for_split`'
-            )
-        pixeltable_schema[column_name_for_split] = ts.StringType(nullable=True)
-
-    for field, column_type in pixeltable_schema.items():
-        if column_type is None:
-            raise excs.Error(f'Could not infer pixeltable type for feature `{field}` in huggingface dataset')
-
-    if isinstance(dataset, datasets.Dataset):
-        # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
-        raw_name = dataset.split._name
-        split_name = raw_name.split('[')[0] if raw_name is not None else None
-        dataset_dict = {split_name: dataset}
-    elif isinstance(dataset, datasets.DatasetDict):
-        dataset_dict = dataset
-    else:
-        raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
-
-    # extract all class labels from the dataset to translate category ints to strings
-    hf_schema = _get_hf_schema(dataset)
-    categorical_features = {
-        feature_name: feature_type.names
-        for (feature_name, feature_type) in hf_schema.items()
-        if isinstance(feature_type, datasets.ClassLabel)
-    }
-
-    try:
-        # random tmp name
-        tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
-        tab = pxt.create_table(tmp_name, pixeltable_schema, **kwargs)
-
-        def _translate_row(row: dict[str, Any], split_name: str) -> dict[str, Any]:
-            output_row = row.copy()
-            # map all class labels to strings
-            for field, values in categorical_features.items():
-                output_row[field] = values[row[field]]
-            # add split name to row
-            if column_name_for_split is not None:
-                output_row[column_name_for_split] = split_name
-            return output_row
-
-        for split_name, split_dataset in dataset_dict.items():
-            num_batches = split_dataset.size_in_bytes / _K_BATCH_SIZE_BYTES
-            tuples_per_batch = math.ceil(split_dataset.num_rows / num_batches)
-            assert tuples_per_batch > 0
-
-            batch = []
-            for row in split_dataset:
-                batch.append(_translate_row(row, split_name))
-                if len(batch) >= tuples_per_batch:
-                    tab.insert(batch)
-                    batch = []
-            # last batch
-            if len(batch) > 0:
-                tab.insert(batch)
-
-    except Exception as e:
-        _logger.error(f'Error while inserting dataset into table: {tmp_name}')
-        raise e
-
-    pxt.move(tmp_name, table_path)
-    return pxt.get_table(table_path)
+    return pxt.create_table(
+        table_path, source=dataset, schema_overrides=schema_overrides, primary_key=primary_key, extra_args=kwargs
+    )