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

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)
+        df = 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])
@@ -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,7 +580,7 @@ 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:
+            if local_annotations_column not in t._get_schema():
                 t.add_columns({local_annotations_column: ts.Json})
 
         resolved_col_mapping = cls.validate_columns(
@@ -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,5 @@
 import os
-from typing import Any, Optional, Union
+from typing import Any
 
 import numpy as np
 import pandas as pd
@@ -16,8 +16,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 +55,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 +84,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 +132,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 +141,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
 
@@ -191,7 +192,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] = {}

pixeltable/io/parquet.py

@@ -1,16 +1,10 @@
 from __future__ import annotations
 
-import datetime
-import io
 import json
 import logging
 import typing
-from collections import deque
 from pathlib import Path
-from typing import Any, Optional, Union
-
-import numpy as np
-import PIL.Image
+from typing import Any
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
@@ -18,31 +12,13 @@ from pixeltable.catalog import Catalog
 from pixeltable.utils.transactional_directory import transactional_directory
 
 if typing.TYPE_CHECKING:
-    import pyarrow as pa
-
     import pixeltable as pxt
 
 _logger = logging.getLogger('pixeltable')
 
 
-def _write_batch(value_batch: dict[str, deque], schema: pa.Schema, output_path: Path) -> None:
-    import pyarrow as pa
-    from pyarrow import parquet
-
-    pydict = {}
-    for field in schema:
-        if isinstance(field.type, pa.FixedShapeTensorType):
-            stacked_arr = np.stack(value_batch[field.name])
-            pydict[field.name] = pa.FixedShapeTensorArray.from_numpy_ndarray(stacked_arr)
-        else:
-            pydict[field.name] = value_batch[field.name]
-
-    tab = pa.Table.from_pydict(pydict, schema=schema)
-    parquet.write_table(tab, str(output_path))
-
-
 def export_parquet(
-    table_or_df: Union[pxt.Table, pxt.DataFrame],
+    table_or_df: pxt.Table | pxt.DataFrame,
     parquet_path: Path,
     partition_size_bytes: int = 100_000_000,
     inline_images: bool = False,
@@ -63,7 +39,9 @@ def export_parquet(
                         If False, will raise an error if the Dataframe has any image column.
                         Default False.
     """
-    from pixeltable.utils.arrow import to_arrow_schema
+    import pyarrow as pa
+
+    from pixeltable.utils.arrow import to_record_batches
 
     df: pxt.DataFrame
     if isinstance(table_or_df, pxt.catalog.Table):
@@ -71,9 +49,6 @@ def export_parquet(
     else:
         df = table_or_df
 
-    type_dict = {k: v.as_dict() for k, v in df.schema.items()}
-    arrow_schema = to_arrow_schema(df.schema)
-
     if not inline_images and any(col_type.is_image_type() for col_type in df.schema.values()):
         raise excs.Error('Cannot export Dataframe with image columns when inline_images is False')
 
@@ -81,78 +56,23 @@ def export_parquet(
     with transactional_directory(parquet_path) as temp_path:
         # dump metadata json file so we can inspect what was the source of the parquet file later on.
         json.dump(df.as_dict(), (temp_path / '.pixeltable.json').open('w'))
+        type_dict = {k: v.as_dict() for k, v in df.schema.items()}
         json.dump(type_dict, (temp_path / '.pixeltable.column_types.json').open('w'))  # keep type metadata
-
         batch_num = 0
-        current_value_batch: dict[str, deque] = {k: deque() for k in df.schema}
-        current_byte_estimate = 0
-
         with Catalog.get().begin_xact(for_write=False):
-            for data_row in df._exec():
-                for (col_name, col_type), e in zip(df.schema.items(), df._select_list_exprs):
-                    val = data_row[e.slot_idx]
-                    if val is None:
-                        current_value_batch[col_name].append(val)
-                        continue
-
-                    assert val is not None
-                    if col_type.is_image_type():
-                        # images get inlined into the parquet file
-                        if data_row.file_paths is not None and data_row.file_paths[e.slot_idx] is not None:
-                            # if there is a file, read directly to preserve information
-                            with open(data_row.file_paths[e.slot_idx], 'rb') as f:
-                                val = f.read()
-                        elif isinstance(val, PIL.Image.Image):
-                            # if no file available, eg. bc it is computed, convert to png
-                            buf = io.BytesIO()
-                            val.save(buf, format='PNG')
-                            val = buf.getvalue()
-                        else:
-                            raise excs.Error(f'unknown image type {type(val)}')
-                        length = len(val)
-                    elif col_type.is_string_type():
-                        length = len(val)
-                    elif col_type.is_video_type() or col_type.is_audio_type():
-                        if data_row.file_paths is not None and data_row.file_paths[e.slot_idx] is not None:
-                            val = data_row.file_paths[e.slot_idx]
-                        else:
-                            raise excs.Error(f'unknown audio/video type {type(val)}')
-                        length = len(val)
-                    elif col_type.is_json_type():
-                        val = json.dumps(val)
-                        length = len(val)
-                    elif col_type.is_array_type():
-                        length = val.nbytes
-                    elif col_type.is_int_type() or col_type.is_float_type():
-                        length = 8
-                    elif col_type.is_bool_type():
-                        length = 1
-                    elif col_type.is_date_type():
-                        length = 4
-                    elif col_type.is_timestamp_type():
-                        val = val.astimezone(datetime.timezone.utc)
-                        length = 8
-                    else:
-                        raise excs.Error(f'unknown type {col_type} for {col_name}')
-
-                    current_value_batch[col_name].append(val)
-                    current_byte_estimate += length
-                if current_byte_estimate > partition_size_bytes:
-                    assert batch_num < 100_000, 'wrote too many parquet files, unclear ordering'
-                    _write_batch(current_value_batch, arrow_schema, temp_path / f'part-{batch_num:05d}.parquet')
-                    batch_num += 1
-                    current_value_batch = {k: deque() for k in df.schema}
-                    current_byte_estimate = 0
-
-            _write_batch(current_value_batch, arrow_schema, temp_path / f'part-{batch_num:05d}.parquet')
+            for record_batch in to_record_batches(df, partition_size_bytes):
+                output_path = temp_path / f'part-{batch_num:05d}.parquet'
+                arrow_tbl = pa.Table.from_batches([record_batch])
+                pa.parquet.write_table(arrow_tbl, str(output_path))
+                batch_num += 1
 
 
 def import_parquet(
     table: str,
     *,
     parquet_path: str,
-    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:
     """Creates a new base table from a Parquet file or set of files. Requires pyarrow to be installed.