pixeltable 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

pixeltable/plan.py

@@ -76,7 +76,8 @@ class Analyzer:
                             f'order_by()'))
                     self.similarity_clause = similarity_clauses[0]
                     img_col = self.similarity_clause.img_col_ref.col
-                    if not img_col.is_indexed:
+                    indexed_col_ids = {info.col.id for info in tbl.tbl_version.idxs_by_name.values()}
+                    if img_col.id not in indexed_col_ids:
                         raise excs.Error(f'nearest() not available for unindexed column {img_col.name}')
 
         # all exprs that are evaluated in Python; not executable
@@ -220,18 +221,11 @@ class Planner:
     ) -> exec.ExecNode:
         """Creates a plan for TableVersion.insert()"""
         assert not tbl.is_view()
-        # things we need to materialize:
-        # 1. stored_cols: all cols we need to store, incl computed cols (and indices)
+        # stored_cols: all cols we need to store, incl computed cols (and indices)
         stored_cols = [c for c in tbl.cols if c.is_stored]
         assert len(stored_cols) > 0
-        # 2. values to insert into indices
-        indexed_cols = [c for c in tbl.cols if c.is_indexed]
-        index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if len(indexed_cols) > 0:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(c, openai_clip) for c in tbl.cols if c.is_indexed]
 
-        row_builder = exprs.RowBuilder([], stored_cols, index_info, [])
+        row_builder = exprs.RowBuilder([], stored_cols, [])
 
         # create InMemoryDataNode for 'rows'
         stored_col_info = row_builder.output_slot_idxs()
@@ -260,7 +254,7 @@ class Planner:
     @classmethod
     def create_update_plan(
             cls, tbl: catalog.TableVersionPath,
-            update_targets: List[Tuple[catalog.Column, exprs.Expr]],
+            update_targets: dict[catalog.Column, exprs.Expr],
             recompute_targets: List[catalog.Column],
             where_clause: Optional[exprs.Predicate], cascade: bool
     ) -> Tuple[exec.ExecNode, List[str], List[catalog.Column]]:
@@ -279,7 +273,7 @@ class Planner:
         # retrieve all stored cols and all target exprs
         assert isinstance(tbl, catalog.TableVersionPath)
         target = tbl.tbl_version  # the one we need to update
-        updated_cols = [col for col, _ in update_targets]
+        updated_cols = list(update_targets.keys())
         if len(recompute_targets) > 0:
             recomputed_cols = recompute_targets.copy()
         else:
@@ -291,12 +285,12 @@ class Planner:
             col for col in target.cols if col.is_stored and not col in updated_cols and not col in recomputed_base_cols
         ]
         select_list = [exprs.ColumnRef(col) for col in copied_cols]
-        select_list.extend([expr for _, expr in update_targets])
+        select_list.extend(update_targets.values())
 
         recomputed_exprs = \
             [c.value_expr.copy().resolve_computed_cols(resolve_cols=recomputed_base_cols) for c in recomputed_base_cols]
         # recomputed cols reference the new values of the updated cols
-        for col, e in update_targets:
+        for col, e in update_targets.items():
             exprs.Expr.list_substitute(recomputed_exprs, exprs.ColumnRef(col), e)
         select_list.extend(recomputed_exprs)
 
@@ -375,16 +369,10 @@ class Planner:
         #   the store
         target = view.tbl_version  # the one we need to populate
         stored_cols = [c for c in target.cols if c.is_stored and (c.is_computed or target.is_iterator_column(c))]
-        # 2. index values
-        indexed_cols = [c for c in target.cols if c.is_indexed]
-        index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if len(indexed_cols) > 0:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(c, openai_clip) for c in target.cols if c.is_indexed]
-        # 3. for component views: iterator args
+        # 2. for component views: iterator args
         iterator_args = [target.iterator_args] if target.iterator_args is not None else []
 
-        row_builder = exprs.RowBuilder(iterator_args, stored_cols, index_info, [])
+        row_builder = exprs.RowBuilder(iterator_args, stored_cols, [])
 
         # execution plan:
         # 1. materialize exprs computed from the base that are needed for stored view columns
@@ -548,7 +536,7 @@ class Planner:
         analyzer = Analyzer(
             tbl, select_list, where_clause=where_clause, group_by_clause=group_by_clause,
             order_by_clause=order_by_clause)
-        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], [], analyzer.sql_exprs)
+        row_builder = exprs.RowBuilder(analyzer.all_exprs, [], analyzer.sql_exprs)
 
         analyzer.finalize(row_builder)
         # select_list: we need to materialize everything that's been collected
@@ -627,21 +615,15 @@ class Planner:
     @classmethod
     def create_add_column_plan(
             cls, tbl: catalog.TableVersionPath, col: catalog.Column
-    ) -> Tuple[exec.ExecNode, Optional[int], Optional[int]]:
+    ) -> Tuple[exec.ExecNode, Optional[int]]:
         """Creates a plan for InsertableTable.add_column()
         Returns:
             plan: the plan to execute
-            ctx: the context to use for the plan
             value_expr slot idx for the plan output (for computed cols)
-            embedding slot idx for the plan output (for indexed image cols)
         """
         assert isinstance(tbl, catalog.TableVersionPath)
         index_info: List[Tuple[catalog.Column, func.Function]] = []
-        if col.is_indexed:
-            from pixeltable.functions.nos.image_embedding import openai_clip
-            index_info = [(col, openai_clip)]
-        row_builder = exprs.RowBuilder(
-            output_exprs=[], columns=[col], indices=index_info, input_exprs=[])
+        row_builder = exprs.RowBuilder(output_exprs=[], columns=[col], input_exprs=[])
         analyzer = Analyzer(tbl, row_builder.default_eval_ctx.target_exprs)
         plan = cls._create_query_plan(tbl, row_builder=row_builder, analyzer=analyzer, with_pk=True)
         plan.ctx.batch_size = 16
@@ -651,6 +633,5 @@ class Planner:
         # we want to flush images
         if col.is_computed and col.is_stored and col.col_type.is_image_type():
             plan.set_stored_img_cols(row_builder.output_slot_idxs())
-        value_expr_slot_idx: Optional[int] = row_builder.output_slot_idxs()[0].slot_idx if col.is_computed else None
-        embedding_slot_idx: Optional[int] = row_builder.index_slot_idxs()[0].slot_idx if col.is_indexed else None
-        return plan, value_expr_slot_idx, embedding_slot_idx
+        value_expr_slot_idx = row_builder.output_slot_idxs()[0].slot_idx if col.is_computed else None
+        return plan, value_expr_slot_idx

pixeltable/store.py

@@ -38,7 +38,7 @@ class StoreBase:
         self.tbl_version = tbl_version
         self.sa_md = sql.MetaData()
         self.sa_tbl: Optional[sql.Table] = None
-        self._create_sa_tbl()
+        self.create_sa_tbl()
 
     def pk_columns(self) -> List[sql.Column]:
         return self._pk_columns
@@ -62,7 +62,7 @@ class StoreBase:
         return [*rowid_cols, self.v_min_col, self.v_max_col]
 
 
-    def _create_sa_tbl(self) -> None:
+    def create_sa_tbl(self) -> None:
         """Create self.sa_tbl from self.tbl_version."""
         system_cols = self._create_system_columns()
         all_cols = system_cols.copy()
@@ -76,9 +76,6 @@ class StoreBase:
                 all_cols.append(col.sa_errormsg_col)
                 all_cols.append(col.sa_errortype_col)
 
-            if col.is_indexed:
-                all_cols.append(col.sa_idx_col)
-
             # we create an index for:
             # - scalar columns (except for strings, because long strings can't be used for B-tree indices)
             # - non-computed video and image columns (they will contain external paths/urls that users might want to
@@ -145,8 +142,8 @@ class StoreBase:
         """Move tmp media files that we generated to a permanent location"""
         for c in media_cols:
             for table_row in table_rows:
-                file_url = table_row[c.storage_name()]
-                table_row[c.storage_name()] = self._move_tmp_media_file(file_url, c, v_min)
+                file_url = table_row[c.store_name()]
+                table_row[c.store_name()] = self._move_tmp_media_file(file_url, c, v_min)
 
     def _create_table_row(
             self, input_row: exprs.DataRow, row_builder: exprs.RowBuilder, media_cols: List[catalog.Column],
@@ -168,16 +165,19 @@ class StoreBase:
 
         return table_row, num_excs
 
-    def count(self) -> None:
+    def count(self, conn: Optional[sql.engine.Connection] = None) -> int:
         """Return the number of rows visible in self.tbl_version"""
         stmt = sql.select(sql.func.count('*'))\
             .select_from(self.sa_tbl)\
             .where(self.v_min_col <= self.tbl_version.version)\
             .where(self.v_max_col > self.tbl_version.version)
-        with env.Env.get().engine.begin() as conn:
+        if conn is None:
+            with env.Env.get().engine.connect() as conn:
+                result = conn.execute(stmt).scalar_one()
+        else:
             result = conn.execute(stmt).scalar_one()
-            assert isinstance(result, int)
-            return result
+        assert isinstance(result, int)
+        return result
 
     def create(self, conn: sql.engine.Connection) -> None:
         self.sa_md.create_all(bind=conn)
@@ -193,38 +193,35 @@ class StoreBase:
         message).
         """
         assert col.is_stored
-        stmt = sql.text(f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.storage_name()} {col.col_type.to_sql()}')
+        col_type_str = col.get_sa_col_type().compile(dialect=conn.dialect)
+        stmt = sql.text(f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.store_name()} {col_type_str} NULL')
         log_stmt(_logger, stmt)
         conn.execute(stmt)
-        added_storage_cols = [col.storage_name()]
+        added_storage_cols = [col.store_name()]
         if col.records_errors:
             # we also need to create the errormsg and errortype storage cols
             stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errormsg_storage_name()} {StringType().to_sql()} DEFAULT NULL')
+                    f'ADD COLUMN {col.errormsg_store_name()} {StringType().to_sql()} DEFAULT NULL')
             conn.execute(sql.text(stmt))
             stmt = (f'ALTER TABLE {self._storage_name()} '
-                    f'ADD COLUMN {col.errortype_storage_name()} {StringType().to_sql()} DEFAULT NULL')
+                    f'ADD COLUMN {col.errortype_store_name()} {StringType().to_sql()} DEFAULT NULL')
             conn.execute(sql.text(stmt))
-        added_storage_cols.extend([col.errormsg_storage_name(), col.errortype_storage_name()])
-        self._create_sa_tbl()
+            added_storage_cols.extend([col.errormsg_store_name(), col.errortype_store_name()])
+        self.create_sa_tbl()
         _logger.info(f'Added columns {added_storage_cols} to storage table {self._storage_name()}')
 
-    def drop_column(self, col: Optional[catalog.Column] = None, conn: Optional[sql.engine.Connection] = None) -> None:
-        """Re-create self.sa_tbl and drop column, if one is given"""
-        if col is not None:
-            assert conn is not None
-            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.storage_name()}'
+    def drop_column(self, col: catalog.Column, conn: sql.engine.Connection) -> None:
+        """Execute Alter Table Drop Column statement"""
+        stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.store_name()}'
+        conn.execute(sql.text(stmt))
+        if col.records_errors:
+            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errormsg_store_name()}'
+            conn.execute(sql.text(stmt))
+            stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errortype_store_name()}'
             conn.execute(sql.text(stmt))
-            if col.records_errors:
-                stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errormsg_storage_name()}'
-                conn.execute(sql.text(stmt))
-                stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.errortype_storage_name()}'
-                conn.execute(sql.text(stmt))
-        self._create_sa_tbl()
 
     def load_column(
-            self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, embedding_slot_idx: int,
-            conn: sql.engine.Connection
+            self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, conn: sql.engine.Connection
     ) -> int:
         """Update store column of a computed column with values produced by an execution plan
 
@@ -253,18 +250,11 @@ class StoreBase:
                             col.sa_errormsg_col: error_msg
                         }
                     else:
-                        val = result_row.get_stored_val(value_expr_slot_idx)
+                        val = result_row.get_stored_val(value_expr_slot_idx, col.sa_col.type)
                         if col.col_type.is_media_type():
                             val = self._move_tmp_media_file(val, col, result_row.pk[-1])
                         values_dict = {col.sa_col: val}
 
-                if col.is_indexed:
-                    # TODO: deal with exceptions
-                    assert not result_row.has_exc(embedding_slot_idx)
-                    # don't use get_stored_val() here, we need to pass the ndarray
-                    embedding = result_row[embedding_slot_idx]
-                    values_dict[col.sa_index_col] = embedding
-
                 update_stmt = sql.update(self.sa_tbl).values(values_dict)
                 for pk_col, pk_val in zip(self.pk_columns(), result_row.pk):
                     update_stmt = update_stmt.where(pk_col == pk_val)
@@ -337,6 +327,7 @@ class StoreBase:
             self, current_version: int, base_versions: List[Optional[int]], match_on_vmin: bool,
             where_clause: Optional[sql.ClauseElement], conn: sql.engine.Connection) -> int:
         """Mark rows as deleted that are live and were created prior to current_version.
+        Also: populate the undo columns
         Args:
             base_versions: if non-None, join only to base rows that were created at that version,
                 otherwise join to rows that are live in the base's current version (which is distinct from the
@@ -354,8 +345,14 @@ class StoreBase:
         rowid_join_clause = self._rowid_join_predicate()
         base_versions_clause = sql.true() if len(base_versions) == 0 \
             else self.base._versions_clause(base_versions, match_on_vmin)
+        set_clause = {self.v_max_col: current_version}
+        for index_info in self.tbl_version.idxs_by_name.values():
+            # copy value column to undo column
+            set_clause[index_info.undo_col.sa_col] = index_info.val_col.sa_col
+            # set value column to NULL
+            set_clause[index_info.val_col.sa_col] = None
         stmt = sql.update(self.sa_tbl) \
-            .values({self.v_max_col: current_version}) \
+            .values(set_clause) \
             .where(where_clause) \
             .where(rowid_join_clause) \
             .where(base_versions_clause)
@@ -416,8 +413,8 @@ class StoreComponentView(StoreView):
         self.rowid_cols.append(self.pos_col)
         return self.rowid_cols
 
-    def _create_sa_tbl(self) -> None:
-        super()._create_sa_tbl()
+    def create_sa_tbl(self) -> None:
+        super().create_sa_tbl()
         # we need to fix up the 'pos' column in TableVersion
         self.tbl_version.cols_by_name['pos'].sa_col = self.pos_col
 

pixeltable/tests/conftest.py

@@ -6,11 +6,12 @@ from typing import List
 
 import numpy as np
 import pytest
+import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.catalog as catalog
 from pixeltable import exprs
-from pixeltable import functions as ptf
+import pixeltable.functions as pxtf
 from pixeltable.exprs import RELATIVE_PATH_ROOT as R
 from pixeltable.metadata import SystemInfo, create_system_info
 from pixeltable.metadata.schema import TableSchemaVersion, TableVersion, Table, Function, Dir
@@ -120,8 +121,7 @@ def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
         t.c1.apply(json.loads),
         t.c8.errortype,
         t.c8.errormsg,
-        ptf.sum(t.c2, group_by=t.c4, order_by=t.c3),
-        #test_stored_fn(t.c2),
+        pxtf.sum(t.c2, group_by=t.c4, order_by=t.c3),
     ]
 
 @pytest.fixture(scope='function')
@@ -153,17 +153,11 @@ def img_tbl_exprs(img_tbl: catalog.Table) -> List[exprs.Expr]:
         img_t.img.localpath,
     ]
 
-# TODO: why does this not work with a session scope? (some user tables don't get created with create_all())
-#@pytest.fixture(scope='session')
-#def indexed_img_tbl(init_env: None) -> catalog.Table:
-#    cl = pxt.Client()
-#    db = cl.create_db('test_indexed')
 @pytest.fixture(scope='function')
-def indexed_img_tbl(test_client: pxt.Client) -> catalog.Table:
-    skip_test_if_not_installed('nos')
+def small_img_tbl(test_client: pxt.Client) -> catalog.Table:
     cl = test_client
     schema = {
-        'img': { 'type': ImageType(nullable=False), 'indexed': True },
+        'img': ImageType(nullable=False),
         'category': StringType(nullable=False),
         'split': StringType(nullable=False),
     }

pixeltable/tests/ext/test_yolox.py

@@ -0,0 +1,21 @@
+import pixeltable as pxt
+from pixeltable.tests.utils import skip_test_if_not_installed, get_image_files, validate_update_status
+
+
+class TestYolox:
+
+    def test_yolox(self, test_client: pxt.Client):
+        skip_test_if_not_installed('yolox')
+        from pixeltable.ext.functions.yolox import yolox
+        cl = test_client
+        t = cl.create_table('yolox_test', {'image': pxt.ImageType()})
+        t['detect_yolox_tiny'] = yolox(t.image, model_id='yolox_tiny')
+        t['detect_yolox_nano'] = yolox(t.image, model_id='yolox_nano', threshold=0.2)
+        t['yolox_nano_bboxes'] = t.detect_yolox_nano.bboxes
+        images = get_image_files()[:10]
+        validate_update_status(t.insert({'image': image} for image in images), expected_rows=10)
+        rows = t.collect()
+        # Verify correctly formed JSON
+        assert all(list(result.keys()) == ['bboxes', 'labels', 'scores'] for result in rows['detect_yolox_tiny'])
+        # Verify that bboxes are actually present in at least some of the rows.
+        assert any(len(bboxes) > 0 for bboxes in rows['yolox_nano_bboxes'])

pixeltable/tests/functions/test_fireworks.py

@@ -5,6 +5,7 @@ import pixeltable.exceptions as excs
 from pixeltable.tests.utils import skip_test_if_not_installed, validate_update_status
 
 
+@pytest.mark.remote_api
 class TestFireworks:
 
     def test_fireworks(self, test_client: pxt.Client) -> None:

pixeltable/tests/functions/test_huggingface.py

@@ -123,10 +123,10 @@ class TestHuggingface:
         for idx, model_id in enumerate(model_ids):
             col_name = f'embed_text{idx}'
             t[col_name] = clip_text(t.text, model_id=model_id)
-            assert t.column_types()[col_name] == ArrayType((None,), dtype=FloatType(), nullable=False)
+            assert t.column_types()[col_name].is_array_type()
             col_name = f'embed_img{idx}'
             t[col_name] = clip_image(t.img, model_id=model_id)
-            assert t.column_types()[col_name] == ArrayType((None,), dtype=FloatType(), nullable=False)
+            assert t.column_types()[col_name].is_array_type()
 
         def verify_row(row: Dict[str, Any]) -> None:
             for idx, _ in enumerate(model_ids):

pixeltable/tests/functions/test_openai.py

@@ -6,6 +6,7 @@ from pixeltable.tests.utils import SAMPLE_IMAGE_URL, skip_test_if_not_installed,
 from pixeltable.type_system import StringType, ImageType
 
 
+@pytest.mark.remote_api
 class TestOpenai:
 
     def test_audio(self, test_client: pxt.Client) -> None:
@@ -22,7 +23,7 @@ class TestOpenai:
         ))
         t.add_column(translation=translations(t.speech, model='whisper-1'))
         t.add_column(translation_2=translations(
-            t.speech, model='whisper-1', prompt='Translate the recording from Spanish into English.', temperature=0.7
+            t.speech, model='whisper-1', prompt='Translate the recording from Spanish into English.', temperature=0.05
         ))
         validate_update_status(t.insert([
             {'input': 'I am a banana.'},
@@ -33,8 +34,8 @@ class TestOpenai:
         results = t.collect()
         assert results[0]['transcription']['text'] in ['I am a banana.', "I'm a banana."]
         assert results[0]['transcription_2']['text'] in ['I am a banana.', "I'm a banana."]
-        assert 'easy to translate from Spanish' in results[1]['translation']['text']
-        assert 'easy to translate from Spanish' in results[1]['translation_2']['text']
+        assert 'easy to translate' in results[1]['translation']['text']
+        assert 'easy to translate' in results[1]['translation_2']['text']
 
     def test_chat_completions(self, test_client: pxt.Client) -> None:
         skip_test_if_not_installed('openai')
@@ -132,13 +133,22 @@ class TestOpenai:
         t.add_column(img_2=image_generations(
             t.input, model='dall-e-2', size='512x512', user='pixeltable'
         ))
+        validate_update_status(t.insert(input='A friendly dinosaur playing tennis in a cornfield'), 1)
+        assert t.collect()['img'][0].size == (1024, 1024)
+        assert t.collect()['img_2'][0].size == (512, 512)
+
+    @pytest.mark.skip('Test is expensive and slow')
+    def test_image_generations_dall_e_3(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('openai')
+        TestOpenai.skip_test_if_no_openai_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.openai import image_generations
         # Test dall-e-3 options
         t.add_column(img_3=image_generations(
             t.input, model='dall-e-3', quality='hd', size='1792x1024', style='natural', user='pixeltable'
         ))
         validate_update_status(t.insert(input='A friendly dinosaur playing tennis in a cornfield'), 1)
-        assert t.collect()['img'][0].size == (1024, 1024)
-        assert t.collect()['img_2'][0].size == (512, 512)
         assert t.collect()['img_3'][0].size == (1792, 1024)
 
     # This ensures that the test will be skipped, rather than returning an error, when no API key is

pixeltable/tests/functions/test_together.py

@@ -5,6 +5,7 @@ import pixeltable.exceptions as excs
 from pixeltable.tests.utils import skip_test_if_not_installed, validate_update_status
 
 
+@pytest.mark.remote_api
 class TestTogether:
 
     def test_completions(self, test_client: pxt.Client) -> None:

pixeltable/tests/test_component_view.py

@@ -9,7 +9,7 @@ import pixeltable as pxt
 from pixeltable import exceptions as excs
 from pixeltable.iterators import ComponentIterator
 from pixeltable.iterators.video import FrameIterator
-from pixeltable.tests.utils import assert_resultset_eq, get_test_video_files
+from pixeltable.tests.utils import assert_resultset_eq, get_test_video_files, validate_update_status
 from pixeltable.type_system import IntType, VideoType, JsonType
 
 class ConstantImgIterator(ComponentIterator):
@@ -157,10 +157,19 @@ class TestComponentView:
         assert status.num_excs == 0
         import urllib
         video_url = urllib.parse.urljoin('file:', urllib.request.pathname2url(video_filepaths[0]))
-        status = view_t.update({'annotation': {'a': 1}}, where=view_t.video == video_url)
-        c1 = view_t.where(view_t.annotation != None).count()
-        c2 = view_t.where(view_t.video == video_url).count()
-        assert c1 == c2
+        validate_update_status(
+            view_t.update({'annotation': {'a': 1}}, where=view_t.video == video_url),
+            expected_rows=view_t.where(view_t.video == video_url).count())
+        assert view_t.where(view_t.annotation != None).count() == view_t.where(view_t.video == video_url).count()
+
+        # batch update with _rowid works
+        validate_update_status(
+            view_t.batch_update(
+                [{'annotation': {'a': 1}, '_rowid': (1, 0)}, {'annotation': {'a': 1}, '_rowid': (1, 1)}]),
+            expected_rows=2)
+        with pytest.raises(AssertionError):
+            # malformed _rowid
+            view_t.batch_update([{'annotation': {'a': 1}, '_rowid': (1,)}])
 
         with pytest.raises(excs.Error) as excinfo:
             _ = cl.create_view(

pixeltable/tests/test_dataframe.py

@@ -16,6 +16,22 @@ from pixeltable.tests.utils import get_video_files, get_audio_files, skip_test_i
 
 
 class TestDataFrame:
+
+    @pxt.udf(return_type=pxt.JsonType(nullable=False), param_types=[pxt.JsonType(nullable=False)])
+    def yolo_to_coco(detections):
+        bboxes, labels = detections['bboxes'], detections['labels']
+        num_annotations = len(detections['bboxes'])
+        assert num_annotations == len(detections['labels'])
+        result = []
+        for i in range(num_annotations):
+            bbox = bboxes[i]
+            ann = {
+                'bbox': [round(bbox[0]), round(bbox[1]), round(bbox[2] - bbox[0]), round(bbox[3] - bbox[1])],
+                'category': labels[i],
+            }
+            result.append(ann)
+        return result
+
     def test_select_where(self, test_tbl: catalog.Table) -> None:
         t = test_tbl
         res1 = t[t.c1, t.c2, t.c3].show(0)
@@ -156,7 +172,7 @@ class TestDataFrame:
         _ = df.__repr__()
         _ = df._repr_html_()
 
-    def test_count(self, test_tbl: catalog.Table, indexed_img_tbl: catalog.Table) -> None:
+    def test_count(self, test_tbl: catalog.Table, small_img_tbl) -> None:
         skip_test_if_not_installed('nos')
         t = test_tbl
         cnt = t.count()
@@ -166,7 +182,7 @@ class TestDataFrame:
         assert cnt == 10
 
         # count() doesn't work with similarity search
-        t = indexed_img_tbl
+        t = small_img_tbl
         probe = t.select(t.img).show(1)
         img = probe[0, 0]
         with pytest.raises(excs.Error):
@@ -397,22 +413,7 @@ class TestDataFrame:
         view_t.add_column(detections=yolox_medium(view_t.frame))
         base_t.insert(video=get_video_files()[0])
 
-        @pxt.udf(return_type=pxt.JsonType(nullable=False), param_types=[pxt.JsonType(nullable=False)])
-        def yolo_to_coco(detections):
-            bboxes, labels = detections['bboxes'], detections['labels']
-            num_annotations = len(detections['bboxes'])
-            assert num_annotations == len(detections['labels'])
-            result = []
-            for i in range(num_annotations):
-                bbox = bboxes[i]
-                ann = {
-                    'bbox': [round(bbox[0]), round(bbox[1]), round(bbox[2] - bbox[0]), round(bbox[3] - bbox[1])],
-                    'category': labels[i],
-                }
-                result.append(ann)
-            return result
-
-        query = view_t.select({'image': view_t.frame, 'annotations': yolo_to_coco(view_t.detections)})
+        query = view_t.select({'image': view_t.frame, 'annotations': self.yolo_to_coco(view_t.detections)})
         path = query.to_coco_dataset()
         # we get a valid COCO dataset
         coco_ds = COCO(path)