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

pixeltable/tests/conftest.py

@@ -1,119 +1,177 @@
-import numpy as np
-import pandas as pd
-import datetime
+import json
+import logging
+import os
+import pathlib
+from typing import List
 
+import numpy as np
 import pytest
 
-import pixeltable as pt
+import pixeltable as pxt
 import pixeltable.catalog as catalog
-from pixeltable.type_system import \
-    StringType, IntType, FloatType, BoolType, TimestampType, ImageType, JsonType
-from pixeltable.tests.utils import read_data_file, make_tbl, create_table_data
+from pixeltable import exprs
+from pixeltable import functions as ptf
+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
+from pixeltable.tests.utils import read_data_file, create_test_tbl, create_all_datatypes_tbl, skip_test_if_not_installed
+from pixeltable.type_system import StringType, ImageType, FloatType
 
 
 @pytest.fixture(scope='session')
-def init_db(tmp_path_factory) -> None:
+def init_env(tmp_path_factory) -> None:
     from pixeltable.env import Env
+    # set the relevant env vars for Client() to connect to the test db
+
+    shared_home = pathlib.Path(os.environ.get('PIXELTABLE_HOME', str(pathlib.Path.home() / '.pixeltable')))
+    home_dir = str(tmp_path_factory.mktemp('base') / '.pixeltable')
+    os.environ['PIXELTABLE_HOME'] = home_dir
+    os.environ['PIXELTABLE_CONFIG'] = str(shared_home / 'config.yaml')
+    test_db = 'test'
+    os.environ['PIXELTABLE_DB'] = test_db
+    os.environ['PIXELTABLE_PGDATA'] = str(shared_home / 'pgdata')
+
+    # ensure this home dir exits
+    shared_home.mkdir(parents=True, exist_ok=True)
     # this also runs create_all()
-    db_name = 'test'
-    Env.get().set_up(tmp_path_factory.mktemp('base'), db_name=db_name, echo=True)
+    Env.get().set_up(echo=True)
     yield
     # leave db in place for debugging purposes
 
+@pytest.fixture(scope='function')
+def test_client(init_env) -> pxt.Client:
+    # Clean the DB *before* instantiating a client object. This is because some tests
+    # (such as test_migration.py) may leave the DB in a broken state, from which the
+    # client is uninstantiable.
+    clean_db()
+    cl = pxt.Client(reload=True)
+    cl.logging(level=logging.DEBUG, to_stdout=True)
+    yield cl
+
+
+def clean_db(restore_tables: bool = True) -> None:
+    from pixeltable.env import Env
+    # The logic from Client.reset_catalog() has been moved here, so that it
+    # does not rely on instantiating a Client object. As before, UUID-named data tables will
+    # not be cleaned. If in the future it is desirable to clean out data tables as well,
+    # the commented lines may be used to drop ALL tables from the test db.
+    # sql_md = declarative_base().metadata
+    # sql_md.reflect(Env.get().engine)
+    # sql_md.drop_all(bind=Env.get().engine)
+    engine = Env.get().engine
+    SystemInfo.__table__.drop(engine, checkfirst=True)
+    TableSchemaVersion.__table__.drop(engine, checkfirst=True)
+    TableVersion.__table__.drop(engine, checkfirst=True)
+    Table.__table__.drop(engine, checkfirst=True)
+    Function.__table__.drop(engine, checkfirst=True)
+    Dir.__table__.drop(engine, checkfirst=True)
+    if restore_tables:
+        Dir.__table__.create(engine)
+        Function.__table__.create(engine)
+        Table.__table__.create(engine)
+        TableVersion.__table__.create(engine)
+        TableSchemaVersion.__table__.create(engine)
+        SystemInfo.__table__.create(engine)
+        create_system_info(engine)
+
 
 @pytest.fixture(scope='function')
-def test_db(init_db: None) -> pt.Db:
-    cl = pt.Client()
-    db = cl.create_db(f'test')
-    yield db
-    cl.drop_db(db.name, force=True)
+def test_tbl(test_client: pxt.Client) -> catalog.Table:
+    return create_test_tbl(test_client)
 
+# @pytest.fixture(scope='function')
+# def test_stored_fn(test_client: pxt.Client) -> pxt.Function:
+#     @pxt.udf(return_type=pxt.IntType(), param_types=[pxt.IntType()])
+#     def test_fn(x):
+#         return x + 1
+#     test_client.create_function('test_fn', test_fn)
+#     return test_fn
 
 @pytest.fixture(scope='function')
-def test_tbl(test_db: pt.Db) -> catalog.Table:
-    cols = [
-        catalog.Column('c1', StringType(), nullable=False),
-        catalog.Column('c2', IntType(), nullable=False),
-        catalog.Column('c3', FloatType(), nullable=False),
-        catalog.Column('c4', BoolType(), nullable=False),
-        catalog.Column('c5', TimestampType(), nullable=False),
-        catalog.Column('c6', JsonType(), nullable=False),
-        catalog.Column('c7', JsonType(), nullable=False),
+def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
+#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pxt.Function) -> List[exprs.Expr]:
+
+    t = test_tbl
+    return [
+        t.c1,
+        t.c7['*'].f1,
+        exprs.Literal('test'),
+        exprs.InlineDict({
+            'a': t.c1, 'b': t.c6.f1, 'c': 17,
+            'd': exprs.InlineDict({'e': t.c2}),
+            'f': exprs.InlineArray((t.c3, t.c3))
+        }),
+        exprs.InlineArray([[t.c2, t.c2], [t.c2, t.c2]]),
+        t.c2 > 5,
+        t.c2 == None,
+        ~(t.c2 > 5),
+        (t.c2 > 5) & (t.c1 == 'test'),
+        (t.c2 > 5) | (t.c1 == 'test'),
+        t.c7['*'].f5 >> [R[3], R[2], R[1], R[0]],
+        t.c8[0, 1:],
+        t.c2.astype(FloatType()),
+        (t.c2 + 1).astype(FloatType()),
+        t.c2.apply(str),
+        (t.c2 + 1).apply(str),
+        t.c3.apply(str),
+        t.c4.apply(str),
+        t.c5.apply(str),
+        t.c6.apply(str),
+        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),
     ]
-    t = test_db.create_table('test__tbl', cols)
-
-    num_rows = 100
-    d1 = {
-        'f1': 'test string 1',
-        'f2': 1,
-        'f3': 1.0,
-        'f4': True,
-        'f5': [1.0, 2.0, 3.0, 4.0],
-        'f6': {
-            'f7': 'test string 2',
-            'f8': [1.0, 2.0, 3.0, 4.0],
-        },
-    }
-    d2 = [d1, d1]
-
-    c1_data = [f'test string {i}' for i in range(num_rows)]
-    c2_data = [i for i in range(num_rows)]
-    c3_data = [float(i) for i in range(num_rows)]
-    c4_data = [bool(i % 2) for i in range(num_rows)]
-    c5_data = [datetime.datetime.now()] * num_rows
-    c6_data = []
-    for i in range(num_rows):
-        d = {
-            'f1': f'test string {i}',
-            'f2': i,
-            'f3': float(i),
-            'f4': bool(i % 2),
-            'f5': [1.0, 2.0, 3.0, 4.0],
-            'f6': {
-                'f7': 'test string 2',
-                'f8': [1.0, 2.0, 3.0, 4.0],
-            },
-        }
-        c6_data.append(d)
-
-    c7_data = [d2] * num_rows
-    data = {'c1': c1_data, 'c2': c2_data, 'c3': c3_data, 'c4': c4_data, 'c5': c5_data, 'c6': c6_data, 'c7': c7_data}
-    pd_df = pd.DataFrame(data=data)
-    t.insert_pandas(pd_df)
-    return t
 
+@pytest.fixture(scope='function')
+def all_datatypes_tbl(test_client: pxt.Client) -> catalog.Table:
+    return create_all_datatypes_tbl(test_client)
 
 @pytest.fixture(scope='function')
-def img_tbl(test_db: pt.Db) -> catalog.Table:
-    cols = [
-        catalog.Column('img', ImageType(), nullable=False, indexed=False),
-        catalog.Column('category', StringType(), nullable=False),
-        catalog.Column('split', StringType(), nullable=False),
-    ]
+def img_tbl(test_client: pxt.Client) -> catalog.Table:
+    schema = {
+        'img': ImageType(nullable=False),
+        'category': StringType(nullable=False),
+        'split': StringType(nullable=False),
+    }
     # this table is not indexed in order to avoid the cost of computing embeddings
-    tbl = test_db.create_table('test_img_tbl', cols)
-    df = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
-    tbl.insert_pandas(df)
+    tbl = test_client.create_table('test_img_tbl', schema)
+    rows = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
+    tbl.insert(rows)
     return tbl
 
+@pytest.fixture(scope='function')
+def img_tbl_exprs(img_tbl: catalog.Table) -> List[exprs.Expr]:
+    img_t = img_tbl
+    return [
+        img_t.img.width,
+        img_t.img.rotate(90),
+        # we're using a list here, not a tuple; the latter turns into a list during the back/forth conversion
+        img_t.img.rotate(90).resize([224, 224]),
+        img_t.img.fileurl,
+        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_db: None) -> catalog.Table:
-#    cl = pt.Client()
+#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_db: pt.Db) -> catalog.Table:
-    db = test_db
-    cols = [
-        catalog.Column('img', ImageType(), nullable=False, indexed=True),
-        catalog.Column('category', StringType(), nullable=False),
-        catalog.Column('split', StringType(), nullable=False),
-    ]
-    tbl = db.create_table('test_indexed_img_tbl', cols)
-    df = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
-    # select rows randomly in the hope of getting a good sample of the available categories
+def indexed_img_tbl(test_client: pxt.Client) -> catalog.Table:
+    skip_test_if_not_installed('nos')
+    cl = test_client
+    schema = {
+        'img': { 'type': ImageType(nullable=False), 'indexed': True },
+        'category': StringType(nullable=False),
+        'split': StringType(nullable=False),
+    }
+    tbl = cl.create_table('test_indexed_img_tbl', schema)
+    rows = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
+    # select output_rows randomly in the hope of getting a good sample of the available categories
     rng = np.random.default_rng(17)
-    idxs = rng.choice(np.arange(len(df)), size=40, replace=False)
-    tbl.insert_pandas(df.iloc[idxs])
+    idxs = rng.choice(np.arange(len(rows)), size=40, replace=False)
+    rows = [rows[i] for i in idxs]
+    tbl.insert(rows)
     return tbl

pixeltable/tests/functions/test_fireworks.py

@@ -0,0 +1,42 @@
+import pytest
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.tests.utils import skip_test_if_not_installed, validate_update_status
+
+
+class TestFireworks:
+
+    def test_fireworks(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('fireworks')
+        TestFireworks.skip_test_if_no_fireworks_client()
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': pxt.StringType()})
+        from pixeltable.functions.fireworks import chat_completions
+        messages = [{'role': 'user', 'content': t.input}]
+        t['output'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat'
+        )
+        t['output_2'] = chat_completions(
+            messages=messages,
+            model='accounts/fireworks/models/llama-v2-7b-chat',
+            max_tokens=300,
+            top_k=40,
+            top_p=0.9,
+            temperature=0.7
+        )
+        validate_update_status(t.insert(input="How's everything going today?"), 1)
+        results = t.collect()
+        assert len(results['output'][0]['choices'][0]['message']['content']) > 0
+        assert len(results['output_2'][0]['choices'][0]['message']['content']) > 0
+
+    # This ensures that the test will be skipped, rather than returning an error, when no API key is
+    # available (for example, when a PR runs in CI).
+    @staticmethod
+    def skip_test_if_no_fireworks_client() -> None:
+        try:
+            import pixeltable.functions.fireworks
+            _ = pixeltable.functions.fireworks.fireworks_client()
+        except excs.Error as exc:
+            pytest.skip(str(exc))

pixeltable/tests/functions/test_functions.py

@@ -0,0 +1,60 @@
+import pixeltable as pxt
+from pixeltable import catalog
+from pixeltable.functions.pil.image import blend
+from pixeltable.iterators import FrameIterator
+from pixeltable.tests.utils import get_video_files, skip_test_if_not_installed
+from pixeltable.type_system import VideoType, StringType
+
+
+class TestFunctions:
+    def test_pil(self, img_tbl: catalog.Table) -> None:
+        t = img_tbl
+        _ = t[t.img, t.img.rotate(90), blend(t.img, t.img.rotate(90), 0.5)].show()
+
+    def test_eval_detections(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('nos')
+        cl = test_client
+        video_t = cl.create_table('video_tbl', {'video': VideoType()})
+        # create frame view
+        args = {'video': video_t.video, 'fps': 1}
+        v = cl.create_view('test_view', video_t, iterator_class=FrameIterator, iterator_args=args)
+
+        files = get_video_files()
+        video_t.insert(video=files[-1])
+        v.add_column(frame_s=v.frame.resize([640, 480]))
+        from pixeltable.functions.nos.object_detection_2d import yolox_nano, yolox_small, yolox_large
+        v.add_column(detections_a=yolox_nano(v.frame_s))
+        v.add_column(detections_b=yolox_small(v.frame_s))
+        v.add_column(gt=yolox_large(v.frame_s))
+        from pixeltable.functions.eval import eval_detections, mean_ap
+        res = v.select(
+            eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels
+            )).show()
+        v.add_column(
+            eval_a=eval_detections(
+                v.detections_a.bboxes, v.detections_a.labels, v.detections_a.scores, v.gt.bboxes, v.gt.labels))
+        v.add_column(
+            eval_b=eval_detections(
+                v.detections_b.bboxes, v.detections_b.labels, v.detections_b.scores, v.gt.bboxes, v.gt.labels))
+        ap_a = v.select(mean_ap(v.eval_a)).show()[0, 0]
+        ap_b = v.select(mean_ap(v.eval_b)).show()[0, 0]
+        common_classes = set(ap_a.keys()) & set(ap_b.keys())
+
+        ## TODO: following assertion is failing on CI, 
+        # It is not necessarily a bug, as assert codition is not expected to be always true
+        # for k in common_classes:
+        # assert ap_a[k] <= ap_b[k]
+
+    def test_str(self, test_client: pxt.Client) -> None:
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType()})
+        from pixeltable.functions.string import str_format
+        t.add_column(s1=str_format('ABC {0}', t.input))
+        t.add_column(s2=str_format('DEF {this}', this=t.input))
+        t.add_column(s3=str_format('GHI {0} JKL {this}', t.input, this=t.input))
+        status = t.insert(input='MNO')
+        assert status.num_rows == 1
+        assert status.num_excs == 0
+        row = t.head()[0]
+        assert row == {'input': 'MNO', 's1': 'ABC MNO', 's2': 'DEF MNO', 's3': 'GHI MNO JKL MNO'}

pixeltable/tests/functions/test_huggingface.py

@@ -0,0 +1,158 @@
+from typing import Dict, Any
+
+import pytest
+
+import pixeltable as pxt
+from pixeltable.tests.utils import skip_test_if_not_installed, get_sentences, get_image_files, \
+    SAMPLE_IMAGE_URL
+from pixeltable.type_system import StringType, JsonType, ImageType, BoolType, FloatType, ArrayType
+
+
+class TestHuggingface:
+
+    def test_hf_function(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('sentence_transformers')
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType(), 'bool_col': BoolType()})
+        from pixeltable.functions.huggingface import sentence_transformer
+        model_id = 'intfloat/e5-large-v2'
+        t.add_column(e5=sentence_transformer(t.input, model_id=model_id))
+        sents = get_sentences()
+        status = t.insert({'input': s, 'bool_col': True} for s in sents)
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+
+        # verify handling of constant params
+        with pytest.raises(ValueError) as exc_info:
+            t.add_column(e5_2=sentence_transformer(t.input, model_id=t.input))
+        assert ': parameter model_id must be a constant value' in str(exc_info.value)
+        with pytest.raises(ValueError) as exc_info:
+            t.add_column(e5_2=sentence_transformer(t.input, model_id=model_id, normalize_embeddings=t.bool_col))
+        assert ': parameter normalize_embeddings must be a constant value' in str(exc_info.value)
+
+        # make sure this doesn't cause an exception
+        # TODO: is there some way to capture the output?
+        t.describe()
+
+    def test_sentence_transformer(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('sentence_transformers')
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType(), 'input_list': JsonType()})
+        sents = get_sentences(10)
+        status = t.insert({'input': s, 'input_list': sents} for s in sents)
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+
+        # run multiple models one at a time in order to exercise batching
+        from pixeltable.functions.huggingface import sentence_transformer, sentence_transformer_list
+        model_ids = ['sentence-transformers/all-mpnet-base-v2', 'BAAI/bge-reranker-base']
+        num_dims = [768, 768]
+        for idx, model_id in enumerate(model_ids):
+            col_name = f'embed{idx}'
+            t[col_name] = sentence_transformer(t.input, model_id=model_id, normalize_embeddings=True)
+            assert t.column_types()[col_name] == ArrayType((None,), dtype=FloatType(), nullable=False)
+            list_col_name = f'embed_list{idx}'
+            t[list_col_name] = sentence_transformer_list(t.input_list, model_id=model_id, normalize_embeddings=True)
+            assert t.column_types()[list_col_name] == JsonType()
+
+        def verify_row(row: Dict[str, Any]) -> None:
+            for idx, (_, d) in enumerate(zip(model_ids, num_dims)):
+                assert row[f'embed{idx}'].shape == (d,)
+                assert len(row[f'embed_list{idx}']) == len(sents)
+                assert all(len(v) == d for v in row[f'embed_list{idx}'])
+
+        verify_row(t.tail(1)[0])
+
+        # execution still works after reload
+        cl = pxt.Client(reload=True)
+        t = cl.get_table('test_tbl')
+        status = t.insert({'input': s, 'input_list': sents} for s in sents)
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+        verify_row(t.tail(1)[0])
+
+    def test_cross_encoder(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('sentence_transformers')
+        cl = test_client
+        t = cl.create_table('test_tbl', {'input': StringType(), 'input_list': JsonType()})
+        sents = get_sentences(10)
+        status = t.insert({'input': s, 'input_list': sents} for s in sents)
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+
+        # run multiple models one at a time in order to exercise batching
+        from pixeltable.functions.huggingface import cross_encoder, cross_encoder_list
+        model_ids = ['cross-encoder/ms-marco-MiniLM-L-6-v2', 'cross-encoder/ms-marco-TinyBERT-L-2-v2']
+        for idx, model_id in enumerate(model_ids):
+            col_name = f'embed{idx}'
+            t[col_name] = cross_encoder(t.input, t.input, model_id=model_id)
+            assert t.column_types()[col_name] == FloatType()
+            list_col_name = f'embed_list{idx}'
+            t[list_col_name] = cross_encoder_list(t.input, t.input_list, model_id=model_id)
+            assert t.column_types()[list_col_name] == JsonType()
+
+        def verify_row(row: Dict[str, Any]) -> None:
+            for i in range(len(model_ids)):
+                assert len(row[f'embed_list{idx}']) == len(sents)
+                assert all(isinstance(v, float) for v in row[f'embed_list{idx}'])
+
+        verify_row(t.tail(1)[0])
+
+        # execution still works after reload
+        cl = pxt.Client(reload=True)
+        t = cl.get_table('test_tbl')
+        status = t.insert({'input': s, 'input_list': sents} for s in sents)
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+        verify_row(t.tail(1)[0])
+
+    def test_clip(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('transformers')
+        cl = test_client
+        t = cl.create_table('test_tbl', {'text': StringType(), 'img': ImageType()})
+        num_rows = 10
+        sents = get_sentences(num_rows)
+        imgs = get_image_files()[:num_rows]
+        status = t.insert({'text': text, 'img': img} for text, img in zip(sents, imgs))
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+
+        # run multiple models one at a time in order to exercise batching
+        from pixeltable.functions.huggingface import clip_text, clip_image
+        model_ids = ['openai/clip-vit-base-patch32', 'laion/CLIP-ViT-B-32-laion2B-s34B-b79K']
+        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)
+            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)
+
+        def verify_row(row: Dict[str, Any]) -> None:
+            for idx, _ in enumerate(model_ids):
+                assert row[f'embed_text{idx}'].shape == (512,)
+                assert row[f'embed_img{idx}'].shape == (512,)
+
+        verify_row(t.tail(1)[0])
+
+        # execution still works after reload
+        cl = pxt.Client(reload=True)
+        t = cl.get_table('test_tbl')
+        status = t.insert({'text': text, 'img': img} for text, img in zip(sents, imgs))
+        assert status.num_rows == len(sents)
+        assert status.num_excs == 0
+        verify_row(t.tail(1)[0])
+
+    def test_detr_for_object_detection(self, test_client: pxt.Client) -> None:
+        skip_test_if_not_installed('transformers')
+        cl = test_client
+        t = cl.create_table('test_tbl', {'img': ImageType()})
+        from pixeltable.functions.huggingface import detr_for_object_detection
+        t['detect'] = detr_for_object_detection(t.img, model_id='facebook/detr-resnet-50', threshold=0.8)
+        status = t.insert(img=SAMPLE_IMAGE_URL)
+        assert status.num_rows == 1
+        assert status.num_excs == 0
+        result = t.select(t.detect).collect()[0]['detect']
+        assert 'orange' in result['label_text']
+        assert 'bowl' in result['label_text']
+        assert 'broccoli' in result['label_text']