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

pixeltable/functions/gemini.py

@@ -1,32 +1,35 @@
 """
 Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
 that wrap various endpoints from the Google Gemini API. In order to use them, you must
-first `pip install google-generativeai` and configure your Gemini credentials, as described in
+first `pip install google-genai` and configure your Gemini credentials, as described in
 the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini) tutorial.
 """
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 import pixeltable as pxt
 from pixeltable import env
 
+if TYPE_CHECKING:
+    from google import genai
+
 
 @env.register_client('gemini')
-def _(api_key: str) -> None:
-    import google.generativeai as genai
+def _(api_key: str) -> 'genai.client.Client':
+    from google import genai
 
-    genai.configure(api_key=api_key)
+    return genai.client.Client(api_key=api_key)
 
 
-def _ensure_loaded() -> None:
-    env.Env.get().get_client('gemini')
+def _genai_client() -> 'genai.client.Client':
+    return env.Env.get().get_client('gemini')
 
 
 @pxt.udf(resource_pool='request-rate:gemini')
 async def generate_content(
     contents: str,
     *,
-    model_name: str,
+    model: str,
     candidate_count: Optional[int] = None,
     stop_sequences: Optional[list[str]] = None,
     max_output_tokens: Optional[int] = None,
@@ -48,11 +51,11 @@ async def generate_content(
 
     __Requirements:__
 
-    - `pip install google-generativeai`
+    - `pip install google-genai`
 
     Args:
         contents: The input content to generate from.
-        model_name: The name of the model to use.
+        model: The name of the model to use.
 
     For details on the other parameters, see: <https://ai.google.dev/gemini-api/docs>
 
@@ -63,14 +66,12 @@ async def generate_content(
         Add a computed column that applies the model `gemini-1.5-flash`
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
-        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model_name='gemini-1.5-flash'))
+        >>> tbl.add_computed_column(response=generate_content(tbl.prompt, model='gemini-1.5-flash'))
     """
-    env.Env.get().require_package('google.generativeai')
-    _ensure_loaded()
-    import google.generativeai as genai
+    env.Env.get().require_package('google.genai')
+    from google.genai import types
 
-    model = genai.GenerativeModel(model_name=model_name)
-    gc = genai.GenerationConfig(
+    config = types.GenerateContentConfig(
         candidate_count=candidate_count,
         stop_sequences=stop_sequences,
         max_output_tokens=max_output_tokens,
@@ -82,10 +83,11 @@ async def generate_content(
         presence_penalty=presence_penalty,
         frequency_penalty=frequency_penalty,
     )
-    response = await model.generate_content_async(contents, generation_config=gc)
-    return response.to_dict()
+
+    response = await _genai_client().aio.models.generate_content(model=model, contents=contents, config=config)
+    return response.model_dump()
 
 
 @generate_content.resource_pool
-def _(model_name: str) -> str:
-    return f'request-rate:gemini:{model_name}'
+def _(model: str) -> str:
+    return f'request-rate:gemini:{model}'

pixeltable/functions/globals.py

@@ -49,22 +49,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     allows_window=True,
     # Allow counting non-null values of any type
     # TODO: should we have an "Any" type that can be used here?
-    type_substitutions=tuple(
-        {T: Optional[t]}  # type: ignore[misc]
-        for t in (
-            ts.String,
-            ts.Int,
-            ts.Float,
-            ts.Bool,
-            ts.Timestamp,
-            ts.Array,
-            ts.Json,
-            ts.Image,
-            ts.Video,
-            ts.Audio,
-            ts.Document,
-        )
-    ),
+    type_substitutions=tuple({T: Optional[t]} for t in ts.ALL_PIXELTABLE_TYPES),  # type: ignore[misc]
 )
 class count(func.Aggregator, typing.Generic[T]):
     def __init__(self) -> None:

pixeltable/functions/huggingface.py

@@ -13,6 +13,7 @@ import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.functions.util import normalize_image_mode, resolve_torch_device
@@ -61,14 +62,14 @@ def sentence_transformer(
 
 
 @sentence_transformer.conditional_return_type
-def _(model_id: str) -> pxt.ArrayType:
+def _(model_id: str) -> ts.ArrayType:
     try:
         from sentence_transformers import SentenceTransformer
 
         model = _lookup_model(model_id, SentenceTransformer)
-        return pxt.ArrayType((model.get_sentence_embedding_dimension(),), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
     except ImportError:
-        return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf
@@ -199,14 +200,14 @@ def _(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[pxt.Array[(None,
 
 
 @clip.conditional_return_type
-def _(model_id: str) -> pxt.ArrayType:
+def _(model_id: str) -> ts.ArrayType:
     try:
         from transformers import CLIPModel
 
         model = _lookup_model(model_id, CLIPModel.from_pretrained)
-        return pxt.ArrayType((model.config.projection_dim,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
     except ImportError:
-        return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
 
 
 @pxt.udf(batch_size=4)

pixeltable/functions/image.py

@@ -16,6 +16,7 @@ from typing import Optional
 import PIL.Image
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.exprs import Expr
 from pixeltable.utils.code import local_public_names
 
@@ -88,10 +89,10 @@ def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
 
 
 @convert.conditional_return_type
-def _(self: Expr, mode: str) -> pxt.ColumnType:
+def _(self: Expr, mode: str) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
 
 
 # Image.crop()
@@ -108,14 +109,12 @@ def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Ima
 
 
 @crop.conditional_return_type
-def _(self: Expr, box: tuple[int, int, int, int]) -> pxt.ColumnType:
+def _(self: Expr, box: tuple[int, int, int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
+    assert isinstance(input_type, ts.ImageType)
     if (isinstance(box, (list, tuple))) and len(box) == 4 and all(isinstance(x, int) for x in box):
-        return pxt.ImageType(
-            size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable
-        )
-    return pxt.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+        return ts.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ts.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
 
 
 # Image.getchannel()
@@ -134,10 +133,10 @@ def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
 
 
 @getchannel.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
+def _(self: Expr) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
 
 
 @pxt.udf(is_method=True)
@@ -183,10 +182,10 @@ def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
 
 
 @resize.conditional_return_type
-def _(self: Expr, size: tuple[int, int]) -> pxt.ColumnType:
+def _(self: Expr, size: tuple[int, int]) -> ts.ColumnType:
     input_type = self.col_type
-    assert isinstance(input_type, pxt.ImageType)
-    return pxt.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
 
 
 # Image.rotate()
@@ -237,7 +236,7 @@ def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
 @rotate.conditional_return_type
 @effect_spread.conditional_return_type
 @transpose.conditional_return_type
-def _(self: Expr) -> pxt.ColumnType:
+def _(self: Expr) -> ts.ColumnType:
     return self.col_type
 
 

pixeltable/functions/json.py

@@ -4,9 +4,10 @@ Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
 Example:
 ```python
 import pixeltable as pxt
+import pixeltable.functions as pxtf
 
 t = pxt.get_table(...)
-t.select(pxt.functions.json.make_list()).collect()
+t.select(pxtf.json.make_list(t.json_col)).collect()
 ```
 """
 

pixeltable/functions/math.py

@@ -1,3 +1,15 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for mathematical operations.
+
+Example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.float_col.floor()).collect()
+```
+"""
+
 import builtins
 import math
 from typing import Optional
@@ -10,6 +22,11 @@ from pixeltable.utils.code import local_public_names
 
 @pxt.udf(is_method=True)
 def abs(self: float) -> float:
+    """
+    Return the absolute value of the given number.
+
+    Equivalent to Python [`builtins.abs()`](https://docs.python.org/3/library/functions.html#abs).
+    """
     return builtins.abs(self)
 
 
@@ -20,6 +37,14 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 @pxt.udf(is_method=True)
 def ceil(self: float) -> float:
+    """
+    Return the ceiling of the given number.
+
+    Equivalent to Python [`float(math.ceil(self))`](https://docs.python.org/3/library/math.html#math.ceil) if `self`
+    is finite, or `self` itself if `self` is infinite. (This is slightly different from the default behavior of
+    `math.ceil(self)`, which always returns an `int` and raises an error if `self` is infinite. The behavior in
+    Pixeltable generalizes the Python operator and is chosen to align with the SQL standard.)
+    """
     # This ensures the same behavior as SQL
     if math.isfinite(self):
         return float(math.ceil(self))
@@ -34,6 +59,14 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 @pxt.udf(is_method=True)
 def floor(self: float) -> float:
+    """
+    Return the ceiling of the given number.
+
+    Equivalent to Python [`float(math.floor(self))`](https://docs.python.org/3/library/math.html#math.ceil) if `self`
+    is finite, or `self` itself if `self` is infinite. (This is slightly different from the default behavior of
+    `math.floor(self)`, which always returns an `int` and raises an error if `self` is infinite. The behavior of
+    Pixeltable generalizes the Python operator and is chosen to align with the SQL standard.)
+    """
     # This ensures the same behavior as SQL
     if math.isfinite(self):
         return float(math.floor(self))
@@ -48,6 +81,13 @@ def _(self: sql.ColumnElement) -> sql.ColumnElement:
 
 @pxt.udf(is_method=True)
 def round(self: float, digits: Optional[int] = None) -> float:
+    """
+    Round a number to a given precision in decimal digits.
+
+    Equivalent to Python [`builtins.round(self, digits or 0)`](https://docs.python.org/3/library/functions.html#round).
+    Note that if `digits` is not specified, the behavior matches `builtins.round(self, 0)` rather than
+    `builtins.round(self)`; this ensures that the return type is always `float` (as in SQL) rather than `int`.
+    """
     # Set digits explicitly to 0 to guarantee a return type of float; this ensures the same behavior as SQL
     return builtins.round(self, digits or 0)
 

pixeltable/functions/mistralai.py

@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Optional, TypeVar, Union
 import numpy as np
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable.env import Env, register_client
 from pixeltable.func.signature import Batch
 from pixeltable.utils.code import local_public_names
@@ -176,9 +177,9 @@ async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     dimensions = _embedding_dimensions_cache.get(model)  # `None` if unknown model
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
 _T = TypeVar('_T')

pixeltable/functions/openai.py

@@ -21,6 +21,7 @@ import numpy as np
 import PIL
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable import env, exprs
 from pixeltable.func import Batch, Tools
 from pixeltable.utils.code import local_public_names
@@ -666,13 +667,13 @@ async def embeddings(
 
 
 @embeddings.conditional_return_type
-def _(model: str, dimensions: Optional[int] = None) -> pxt.ArrayType:
+def _(model: str, dimensions: Optional[int] = None) -> ts.ArrayType:
     if dimensions is None:
         if model not in _embedding_dimensions_cache:
             # TODO: find some other way to retrieve a sample
-            return pxt.ArrayType((None,), dtype=pxt.FloatType(), nullable=False)
+            return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
         dimensions = _embedding_dimensions_cache.get(model)
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType(), nullable=False)
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType(), nullable=False)
 
 
 #####################################
@@ -738,17 +739,17 @@ async def image_generations(
 
 
 @image_generations.conditional_return_type
-def _(size: Optional[str] = None) -> pxt.ImageType:
+def _(size: Optional[str] = None) -> ts.ImageType:
     if size is None:
-        return pxt.ImageType(size=(1024, 1024))
+        return ts.ImageType(size=(1024, 1024))
     x_pos = size.find('x')
     if x_pos == -1:
-        return pxt.ImageType()
+        return ts.ImageType()
     try:
         width, height = int(size[:x_pos]), int(size[x_pos + 1 :])
     except ValueError:
-        return pxt.ImageType()
-    return pxt.ImageType(size=(width, height))
+        return ts.ImageType()
+    return ts.ImageType(size=(width, height))
 
 
 #####################################

pixeltable/functions/string.py

@@ -5,10 +5,9 @@ It closely follows the Pandas `pandas.Series.str` API.
 Example:
 ```python
 import pixeltable as pxt
-from pixeltable.functions import string as pxt_str
 
 t = pxt.get_table(...)
-t.select(pxt_str.capitalize(t.str_col)).collect()
+t.select(t.str_col.capitalize()).collect()
 ```
 """
 

pixeltable/functions/together.py

@@ -16,6 +16,7 @@ import tenacity
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.utils.code import local_public_names
@@ -225,12 +226,12 @@ async def embeddings(input: Batch[str], *, model: str) -> Batch[pxt.Array[(None,
 
 
 @embeddings.conditional_return_type
-def _(model: str) -> pxt.ArrayType:
+def _(model: str) -> ts.ArrayType:
     if model not in _embedding_dimensions_cache:
         # TODO: find some other way to retrieve a sample
-        return pxt.ArrayType((None,), dtype=pxt.FloatType())
+        return ts.ArrayType((None,), dtype=ts.FloatType())
     dimensions = _embedding_dimensions_cache[model]
-    return pxt.ArrayType((dimensions,), dtype=pxt.FloatType())
+    return ts.ArrayType((dimensions,), dtype=ts.FloatType())
 
 
 @pxt.udf(resource_pool='request-rate:together:images')

pixeltable/functions/video.py

@@ -4,10 +4,10 @@ Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
 Example:
 ```python
 import pixeltable as pxt
-from pixeltable.functions import video as pxt_video
+import pixeltable.functions as pxtf
 
 t = pxt.get_table(...)
-t.select(pxt_video.extract_audio(t.video_col)).collect()
+t.select(pxtf.video.extract_audio(t.video_col)).collect()
 ```
 """
 

pixeltable/globals.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import logging
 import os
-import urllib.parse
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterable, Iterator, Literal, Optional, Union
 
@@ -372,6 +371,31 @@ def create_snapshot(
     )
 
 
+def create_replica(destination: str, source: Union[str, catalog.Table]) -> Optional[catalog.Table]:
+    """
+    Create a replica of a table. Can be used either to create a remote replica of a local table, or to create a local
+    replica of a remote table. A given table can have at most one replica per Pixeltable instance.
+
+    Args:
+        destination: Path where the replica will be created. Can be either a local path such as `'my_dir.my_table'`, or
+            a remote URI such as `'pxt://username/mydir.my_table'`.
+        source: Path to the source table, or (if the source table is a local table) a handle to the source table.
+    """
+    remote_dest = destination.startswith('pxt://')
+    remote_source = isinstance(source, str) and source.startswith('pxt://')
+    if remote_dest == remote_source:
+        raise excs.Error('Exactly one of `destination` or `source` must be a remote URI.')
+
+    if remote_dest:
+        if isinstance(source, str):
+            source = get_table(source)
+        share.push_replica(destination, source)
+        return None
+    else:
+        assert isinstance(source, str)
+        return share.pull_replica(destination, source)
+
+
 def get_table(path: str) -> catalog.Table:
     """Get a handle to an existing table, view, or snapshot.
 
@@ -470,7 +494,7 @@ def drop_table(
         # if we're dropping a table by handle, we first need to get the current path, then drop the S lock on
         # the Table record, and then get X locks in the correct order (first containing directory, then table)
         with Env.get().begin_xact():
-            tbl_path = table._path()
+            tbl_path = table._path
     else:
         assert isinstance(table, str)
         tbl_path = table
@@ -627,13 +651,6 @@ def _extract_paths(
     return result
 
 
-def publish_snapshot(dest_uri: str, table: catalog.Table) -> None:
-    parsed_uri = urllib.parse.urlparse(dest_uri)
-    if parsed_uri.scheme != 'pxt':
-        raise excs.Error(f'Invalid Pixeltable URI (does not start with pxt://): {dest_uri}')
-    share.publish_snapshot(dest_uri, table)
-
-
 def list_dirs(path: str = '', recursive: bool = True) -> list[str]:
     """List the directories in a directory.
 

pixeltable/io/datarows.py

@@ -3,13 +3,14 @@ from __future__ import annotations
 from typing import Any, Iterable, Optional, Union
 
 import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable import exceptions as excs
 
 
 def _infer_schema_from_rows(
     rows: Iterable[dict[str, Any]], schema_overrides: dict[str, Any], primary_key: list[str]
-) -> dict[str, pxt.ColumnType]:
-    schema: dict[str, pxt.ColumnType] = {}
+) -> dict[str, ts.ColumnType]:
+    schema: dict[str, ts.ColumnType] = {}
     cols_with_nones: set[str] = set()
 
     for n, row in enumerate(rows):
@@ -23,7 +24,7 @@ def _infer_schema_from_rows(
             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=col_name not in primary_key)
+                col_type = ts.ColumnType.infer_literal_type(value, nullable=col_name not in primary_key)
                 if col_type is None:
                     raise excs.Error(
                         f'Could not infer type for column `{col_name}`; the value in row {n} '

pixeltable/io/hf_datasets.py

@@ -31,8 +31,8 @@ _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),
-    'date32': ts.StringType(nullable=True),  # date32 is not supported in pixeltable, use string
-    'date64': ts.StringType(nullable=True),  # date64 is not supported in pixeltable, use string
+    'date32': ts.DateType(nullable=True),
+    'date64': ts.DateType(nullable=True),
 }
 
 

pixeltable/io/label_studio.py

@@ -11,7 +11,7 @@ import label_studio_sdk  # type: ignore[import-untyped]
 import PIL.Image
 from requests.exceptions import HTTPError
 
-import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable import Column, Table, env, exceptions as excs
 from pixeltable.config import Config
 from pixeltable.exprs import ColumnRef, DataRow, Expr
@@ -89,21 +89,21 @@ class LabelStudioProject(Project):
     def __project_config(self) -> '_LabelStudioConfig':
         return self.__parse_project_config(self.project_params['label_config'])
 
-    def get_export_columns(self) -> dict[str, pxt.ColumnType]:
+    def get_export_columns(self) -> dict[str, ts.ColumnType]:
         """
         The data keys and preannotation fields specified in this Label Studio project.
         """
         return self.__project_config.export_columns
 
-    def get_import_columns(self) -> dict[str, pxt.ColumnType]:
+    def get_import_columns(self) -> dict[str, ts.ColumnType]:
         """
         Always contains a single entry:
 
         ```
-        {"annotations": pxt.JsonType(nullable=True)}
+        {"annotations": ts.JsonType(nullable=True)}
         ```
         """
-        return {ANNOTATIONS_COLUMN: pxt.JsonType(nullable=True)}
+        return {ANNOTATIONS_COLUMN: ts.JsonType(nullable=True)}
 
     def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
         _logger.info(
@@ -412,8 +412,8 @@ 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 is not None
-                ancestor = ancestor._base
+                assert ancestor._base_table is not None
+                ancestor = ancestor._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)
@@ -577,10 +577,10 @@ class LabelStudioProject(Project):
             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: pxt.JsonType(nullable=True)})
+                t.add_columns({local_annotations_column: ts.JsonType(nullable=True)})
 
         resolved_col_mapping = cls.validate_columns(
-            t, config.export_columns, {ANNOTATIONS_COLUMN: pxt.JsonType(nullable=True)}, col_mapping
+            t, config.export_columns, {ANNOTATIONS_COLUMN: ts.JsonType(nullable=True)}, col_mapping
         )
 
         # Perform some additional validation
@@ -649,7 +649,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
-    column_type: pxt.ColumnType
+    column_type: ts.ColumnType
 
 
 @dataclass(frozen=True)
@@ -673,18 +673,18 @@ class _LabelStudioConfig:
                 )
 
     @property
-    def export_columns(self) -> dict[str, pxt.ColumnType]:
+    def export_columns(self) -> dict[str, ts.ColumnType]:
         data_key_cols = {key_id: key_info.column_type for key_id, key_info in self.data_keys.items()}
-        rl_cols = {name: pxt.JsonType() for name in self.rectangle_labels}
+        rl_cols = {name: ts.JsonType() for name in self.rectangle_labels}
         return {**data_key_cols, **rl_cols}
 
 
 ANNOTATIONS_COLUMN = 'annotations'
 _PAGE_SIZE = 100  # This is the default used in the LS SDK
 _LS_TAG_MAP = {
-    'header': pxt.StringType(),
-    'text': pxt.StringType(),
-    'image': pxt.ImageType(),
-    'video': pxt.VideoType(),
-    'audio': pxt.AudioType(),
+    'header': ts.StringType(),
+    'text': ts.StringType(),
+    'image': ts.ImageType(),
+    'video': ts.VideoType(),
+    'audio': ts.AudioType(),
 }