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

pixeltable/tests/test_table.py

@@ -8,6 +8,7 @@ import PIL
 import cv2
 import numpy as np
 import pandas as pd
+import pathlib
 import pytest
 
 import pixeltable as pxt
@@ -17,7 +18,7 @@ from pixeltable import exceptions as excs
 from pixeltable.iterators import FrameIterator
 from pixeltable.tests.utils import \
     make_tbl, create_table_data, read_data_file, get_video_files, get_audio_files, get_image_files, get_documents, \
-    assert_resultset_eq
+    assert_resultset_eq, assert_hf_dataset_equal, make_test_arrow_table, validate_update_status
 from pixeltable.tests.utils import skip_test_if_not_installed
 from pixeltable.type_system import \
     StringType, IntType, FloatType, TimestampType, ImageType, VideoType, JsonType, BoolType, ArrayType, AudioType, \
@@ -25,7 +26,6 @@ from pixeltable.type_system import \
 from pixeltable.utils.filecache import FileCache
 from pixeltable.utils.media_store import MediaStore
 
-
 class TestTable:
     # exc for a % 10 == 0
     @pxt.udf(return_type=FloatType(), param_types=[IntType()])
@@ -41,6 +41,21 @@ class TestTable:
     def add1(a: int) -> int:
         return a + 1
 
+    @pxt.uda(
+        update_types=[IntType()], value_type=IntType(), requires_order_by=True,
+        allows_window=True)
+    class window_fn:
+        def __init__(self):
+            pass
+        def update(self, i: int) -> None:
+            pass
+        def value(self) -> int:
+            return 1
+
+    @pxt.expr_udf(param_types=[IntType(nullable=False)])
+    def add1(a: int) -> int:
+        return a + 1
+
     def test_create(self, test_client: pxt.Client) -> None:
         cl = test_client
         cl.create_dir('dir1')
@@ -56,7 +71,7 @@ class TestTable:
         with pytest.raises(excs.Error):
             _ = cl.create_table('1test', schema)
         with pytest.raises(excs.Error):
-            _ = catalog.Column('1c', StringType())
+            _ = cl.create_table('bad name', schema={'c1': StringType()})
         with pytest.raises(excs.Error):
             _ = cl.create_table('test', schema)
         with pytest.raises(excs.Error):
@@ -116,6 +131,100 @@ class TestTable:
         tbl.revert()
         assert tbl.num_retained_versions == num_retained_versions
 
+    def test_import_parquet(self, test_client: pxt.Client, tmp_path: pathlib.Path) -> None:
+        skip_test_if_not_installed('pyarrow')
+        import pyarrow as pa
+        from pixeltable.utils.arrow import iter_tuples
+
+        parquet_dir = tmp_path / 'test_data'
+        parquet_dir.mkdir()
+        make_test_arrow_table(parquet_dir)
+
+        tab = test_client.import_parquet('test_parquet', parquet_path=str(parquet_dir))
+        assert 'test_parquet' in test_client.list_tables()
+        assert tab is not None
+        num_elts = tab.count()
+        arrow_tab: pa.Table = pa.parquet.read_table(str(parquet_dir))
+        assert num_elts == arrow_tab.num_rows
+        assert set(tab.column_names()) == set(arrow_tab.column_names)
+
+        result_set = tab.order_by(tab.c_id).collect()
+        column_types = tab.column_types()
+
+        for tup, arrow_tup in zip(result_set, iter_tuples(arrow_tab)):
+            assert tup['c_id'] == arrow_tup['c_id']
+            for col, val in tup.items():
+                if val is None:
+                    assert arrow_tup[col] is None
+                    continue
+
+                if column_types[col].is_array_type():
+                    assert (val == arrow_tup[col]).all()
+                else:
+                    assert val == arrow_tup[col]
+
+    def test_import_huggingface_dataset(self, test_client: pxt.Client, tmp_path: pathlib.Path) -> None:
+        skip_test_if_not_installed('datasets')
+        import datasets
+
+        test_cases = [
+            # { # includes a timestamp. 20MB for specific slice
+            # Disbled this test case because download is failing, and its not critical.
+            #     'dataset_name': 'c4',
+            #     # see https://huggingface.co/datasets/allenai/c4/blob/main/realnewslike/c4-train.00000-of-00512.json.gz
+            #     'dataset': datasets.load_dataset(
+            #         "allenai/c4",
+            #         data_dir="realnewslike",
+            #         data_files="c4-train.00000-of-00512.json.gz",
+            #         split='train[:1000]',
+            #         cache_dir=tmp_path
+            #     ),
+            # },
+            {  # includes an embedding (array type), common in a few RAG datasets.
+                'dataset_name': 'cohere_wikipedia',
+                'dataset': datasets.load_dataset("Cohere/wikipedia-2023-11-embed-multilingual-v3",
+                                                 data_dir='cr').select_columns(['url', 'title', 'text', 'emb']),
+                # column with name `_id`` is not currently allowed by pixeltable rules,
+                # so filter out that column.
+                # cr subdir has a small number of rows, avoid running out of space in CI runner
+                # see https://huggingface.co/datasets/Cohere/wikipedia-2023-11-embed-multilingual-v3/tree/main/cr
+                'schema_override': {'emb': ArrayType((1024,), dtype=FloatType(), nullable=False)}
+            },
+            # example of dataset dictionary with multiple splits
+            {
+                'dataset_name': 'rotten_tomatoes',
+                'dataset': datasets.load_dataset("rotten_tomatoes"),
+            },
+        ]
+
+        # test a column name for splits other than the default of 'split'
+        split_column_name = 'my_split_col'
+        for rec in test_cases:
+            dataset_name = rec['dataset_name']
+            hf_dataset = rec['dataset']
+
+            tab = test_client.import_huggingface_dataset(
+                dataset_name,
+                hf_dataset,
+                column_name_for_split=split_column_name,
+                schema_override=rec.get('schema_override', None),
+            )
+            if isinstance(hf_dataset, datasets.Dataset):
+                assert_hf_dataset_equal(hf_dataset, tab.df(), split_column_name)
+            elif isinstance(hf_dataset, datasets.DatasetDict):
+                assert tab.count() == sum(hf_dataset.num_rows.values())
+                assert split_column_name in tab.column_names()
+
+                for dataset_name in hf_dataset:
+                    df = tab.where(tab.my_split_col == dataset_name)
+                    assert_hf_dataset_equal(hf_dataset[dataset_name], df, split_column_name)
+            else:
+                assert False
+
+        with pytest.raises(excs.Error) as exc_info:
+            test_client.import_huggingface_dataset('test', {})
+        assert 'type(dataset)' in str(exc_info.value)
+
     def test_image_table(self, test_client: pxt.Client) -> None:
         n_sample_rows = 20
         cl = test_client
@@ -205,10 +314,6 @@ class TestTable:
             cl.create_table('test', {'c1': {'type': StringType(), 'stored': 'true'}})
         assert '"stored" must be a bool' in str(exc_info.value)
 
-        with pytest.raises(excs.Error) as exc_info:
-            cl.create_table('test', {'c1': {'type': StringType(), 'indexed': 'true'}})
-        assert '"indexed" must be a bool' in str(exc_info.value)
-
         with pytest.raises(excs.Error) as exc_info:
             cl.create_table('test', {'c1': StringType()}, primary_key='c2')
         assert 'primary key column c2 not found' in str(exc_info.value).lower()
@@ -407,18 +512,8 @@ class TestTable:
         # a non-materialized column that refers to another non-materialized column
         view.add_column(c4=view.c2.rotate(60), stored=False)
 
-        @pxt.uda(
-            name='window_fn', update_types=[IntType()], value_type=IntType(), requires_order_by = True,
-            allows_window = True)
-        class WindowFnAggregator:
-            def __init__(self):
-                pass
-            def update(self, i: int) -> None:
-                pass
-            def value(self) -> int:
-                return 1
         # cols computed with window functions are stored by default
-        view.add_column(c5=window_fn(view.frame_idx, 1, group_by=view.video))
+        view.add_column(c5=self.window_fn(view.frame_idx, 1, group_by=view.video))
 
         # reload to make sure that metadata gets restored correctly
         cl = pxt.Client(reload=True)
@@ -459,6 +554,23 @@ class TestTable:
         cl.drop_table('test_tbl')
         assert MediaStore.count(view.get_id()) == 0
 
+    def test_insert_nulls(self, test_client: pxt.Client) -> None:
+        cl = test_client
+        schema = {
+            'c1': StringType(nullable=True),
+            'c2': IntType(nullable=True),
+            'c3': FloatType(nullable=True),
+            'c4': BoolType(nullable=True),
+            'c5': ArrayType((2, 3), dtype=IntType(), nullable=True),
+            'c6': JsonType(nullable=True),
+            'c7': ImageType(nullable=True),
+            'c8': VideoType(nullable=True),
+        }
+        t = cl.create_table('test1', schema)
+        status = t.insert(c1='abc')
+        assert status.num_rows == 1
+        assert status.num_excs == 0
+
     def test_insert(self, test_client: pxt.Client) -> None:
         cl = test_client
         schema = {
@@ -533,6 +645,15 @@ class TestTable:
             t.insert(c5=np.ndarray((3, 2)))
         assert 'expected ndarray((2, 3)' in str(exc_info.value)
 
+    def test_insert_string_with_null(self, test_client: pxt.Client) -> None:
+        cl = test_client
+        t = cl.create_table('test', {'c1': StringType()})
+
+        t.insert([{'c1': 'this is a python\x00string'}])
+        assert t.count() == 1
+        for tup in t.df().collect():
+            assert tup['c1'] == 'this is a python string'
+
     def test_query(self, test_client: pxt.Client) -> None:
         skip_test_if_not_installed('boto3')
         cl = test_client
@@ -547,7 +668,63 @@ class TestTable:
         t2 = cl.get_table('test')
         _  = t2.show(n=0)
 
-    def test_update(self, test_tbl: pxt.Table, indexed_img_tbl: pxt.Table) -> None:
+    def test_batch_update(self, test_tbl: pxt.Table) -> None:
+        t = test_tbl
+        validate_update_status(
+            t.batch_update([{'c1': '1', 'c2': 1}, {'c1': '2', 'c2': 2}]),
+            expected_rows=2)
+        assert t.where(t.c2 == 1).collect()[0]['c1'] == '1'
+        assert t.where(t.c2 == 2).collect()[0]['c1'] == '2'
+        validate_update_status(
+            t.batch_update([{'c1': 'one', '_rowid': (1,)}, {'c1': 'two', '_rowid': (2,)}]),
+            expected_rows=2)
+        assert t.where(t.c2 == 1).collect()[0]['c1'] == 'one'
+        assert t.where(t.c2 == 2).collect()[0]['c1'] == 'two'
+
+        cl = pxt.Client()
+        # test composite primary key
+        schema = {'c1': StringType(), 'c2': IntType(), 'c3': FloatType()}
+        t = cl.create_table('composite', schema=schema, primary_key=['c1', 'c2'])
+        rows = [{'c1': str(i), 'c2': i, 'c3': float(i)} for i in range(10)]
+        validate_update_status(t.insert(rows), expected_rows=10)
+
+        validate_update_status(
+            t.batch_update([{'c1': '1', 'c2': 1, 'c3': 2.0}, {'c1': '2', 'c2': 2, 'c3': 3.0}]),
+            expected_rows=2)
+
+        with pytest.raises(excs.Error) as exc_info:
+            # can't mix _rowid with primary key
+            _ = t.batch_update([{'c1': '1', 'c2': 1, 'c3': 2.0, '_rowid': (1,)}])
+        assert 'c1 is a primary key column' in str(exc_info.value).lower()
+
+        with pytest.raises(excs.Error) as exc_info:
+            # bad literal
+            _ = t.batch_update([{'c2': 1, 'c3': 'a'}])
+        assert "'a' is not a valid literal" in str(exc_info.value).lower()
+
+        with pytest.raises(excs.Error) as exc_info:
+            # missing primary key column
+            t.batch_update([{'c1': '1', 'c3': 2.0}])
+        assert 'primary key columns (c2) missing' in str(exc_info.value).lower()
+
+        # table without primary key
+        t2 = cl.create_table('no_pk', schema=schema)
+        validate_update_status(t2.insert(rows), expected_rows=10)
+        with pytest.raises(excs.Error) as exc_info:
+            _ = t2.batch_update([{'c1': '1', 'c2': 1, 'c3': 2.0}])
+        assert 'must have primary key for batch update' in str(exc_info.value).lower()
+
+        # updating with _rowid still works
+        validate_update_status(
+            t2.batch_update([{'c1': 'one', '_rowid': (1,)}, {'c1': 'two', '_rowid': (2,)}]),
+            expected_rows=2)
+        assert t2.where(t2.c2 == 1).collect()[0]['c1'] == 'one'
+        assert t2.where(t2.c2 == 2).collect()[0]['c1'] == 'two'
+        with pytest.raises(AssertionError):
+            # some rows are missing rowids
+            _ = t2.batch_update([{'c1': 'one', '_rowid': (1,)}, {'c1': 'two'}])
+
+    def test_update(self, test_tbl: pxt.Table, small_img_tbl) -> None:
         t = test_tbl
         # update every type with a literal
         test_cases = [
@@ -652,7 +829,7 @@ class TestTable:
             t.update({'c3': 1.0}, where=lambda c2: c2 == 10)
         assert 'Predicate' in str(excinfo.value)
 
-        img_t = indexed_img_tbl
+        img_t = small_img_tbl
 
         # can't update image col
         with pytest.raises(excs.Error) as excinfo:
@@ -682,7 +859,7 @@ class TestTable:
         r2 = t.where(t.c2 < 5).select(t.c3, t.c10, t.d1, t.d2).order_by(t.c2).show(0)
         assert_resultset_eq(r1, r2)
 
-    def test_delete(self, test_tbl: pxt.Table, indexed_img_tbl: pxt.Table) -> None:
+    def test_delete(self, test_tbl: pxt.Table, small_img_tbl) -> None:
         t = test_tbl
 
         cnt = t.where(t.c3 < 10.0).count()
@@ -710,7 +887,7 @@ class TestTable:
             t.delete(where=lambda c2: c2 == 10)
         assert 'Predicate' in str(excinfo.value)
 
-        img_t = indexed_img_tbl
+        img_t = small_img_tbl
         # similarity search is not supported
         with pytest.raises(excs.Error) as excinfo:
             img_t.delete(where=img_t.img.nearest('car'))
@@ -876,6 +1053,10 @@ class TestTable:
         t2.revert()
         assert MediaStore.count(t2.get_id()) == t2.count() * stores_img_col
 
+    @pxt.udf(return_type=ImageType(), param_types=[ImageType()])
+    def img_fn_with_exc(img: PIL.Image.Image) -> PIL.Image.Image:
+        raise RuntimeError
+
     def test_computed_img_cols(self, test_client: pxt.Client) -> None:
         cl = test_client
         schema = {'img': ImageType(nullable=False)}
@@ -893,10 +1074,7 @@ class TestTable:
 
         # computed img col with exceptions
         t = cl.create_table('test3', schema)
-        @pxt.udf(return_type=ImageType(), param_types=[ImageType()])
-        def f(img: PIL.Image.Image) -> PIL.Image.Image:
-            raise RuntimeError
-        t.add_column(c3=f(t.img), stored=True)
+        t.add_column(c3=self.img_fn_with_exc(t.img), stored=True)
         rows = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
         rows = [{'img': r['img']} for r in rows[:20]]
         t.insert(rows, fail_on_exception=False)
@@ -1108,6 +1286,7 @@ class TestTable:
         check_rename(t, 'c1_renamed', 'c1')
 
         # revert() works
+        _ = t.select(t.c1_renamed).collect()
         t.revert()
         _ = t.select(t.c1).collect()
         #check_rename(t, 'c1', 'c1_renamed')

pixeltable/tests/test_types.py

@@ -1,3 +1,7 @@
+import datetime
+from copy import copy
+from typing import List, Dict, Optional
+
 from pixeltable.type_system import \
     ColumnType, StringType, IntType, BoolType, ImageType, InvalidType, FloatType, TimestampType, JsonType, ArrayType
 
@@ -20,3 +24,29 @@ class TestTypes:
             t_serialized = t.serialize()
             t_deserialized = ColumnType.deserialize(t_serialized)
             assert t == t_deserialized
+
+    def test_from_python_type(self) -> None:
+        test_cases = {
+            str: StringType(),
+            int: IntType(),
+            float: FloatType(),
+            bool: BoolType(),
+            datetime.date: TimestampType(),
+            datetime.datetime: TimestampType(),
+            list: JsonType(),
+            dict: JsonType(),
+            list[int]: JsonType(),
+            list[dict[str, int]]: JsonType(),
+            dict[int, str]: JsonType(),
+            dict[dict[str, int], list[int]]: JsonType(),
+            List: JsonType(),
+            Dict: JsonType(),
+            List[int]: JsonType(),
+            List[Dict[str, int]]: JsonType(),
+            Dict[int, str]: JsonType()
+        }
+        for py_type, pxt_type in test_cases.items():
+            assert ColumnType.from_python_type(py_type) == pxt_type
+            opt_pxt_type = copy(pxt_type)
+            opt_pxt_type.nullable = True
+            assert ColumnType.from_python_type(Optional[py_type]) == opt_pxt_type

pixeltable/tests/test_video.py

@@ -107,6 +107,18 @@ class TestVideo:
         base_t.insert({'video': p} for p in video_filepaths)
         _ = view_t[view_t.c1, view_t.c2, view_t.c3, view_t.c4].show(0)
 
+    # window function that simply passes through the frame
+    @pxt.uda(
+        update_types=[ImageType()], value_type=ImageType(),
+        requires_order_by=True, allows_std_agg=False, allows_window=True)
+    class agg_fn:
+        def __init__(self):
+            self.img = None
+        def update(self, frame: PIL.Image.Image) -> None:
+            self.img = frame
+        def value(self) -> PIL.Image.Image:
+            return self.img
+
     def test_make_video(self, test_client: pxt.Client) -> None:
         video_filepaths = get_video_files()
         cl = test_client
@@ -131,29 +143,17 @@ class TestVideo:
                 make_video(view_t.pos, view_t.frame),
                 make_video(view_t.pos - 1, view_t.transformed)).group_by(base_t).show()
 
-        # window function that simply passes through the frame
-        @pxt.uda(
-            update_types=[ImageType()], value_type=ImageType(), name='agg_fn',
-            requires_order_by=True, allows_std_agg=False, allows_window=True)
-        class WindowAgg:
-            def __init__(self):
-                self.img = None
-            def update(self, frame: PIL.Image.Image) -> None:
-                self.img = frame
-            def value(self) -> PIL.Image.Image:
-                return self.img
-
         # make sure it works
-        _ = view_t.select(agg_fn(view_t.pos, view_t.frame, group_by=base_t)).show()
-        status = view_t.add_column(agg=agg_fn(view_t.pos, view_t.frame, group_by=base_t))
+        _ = view_t.select(self.agg_fn(view_t.pos, view_t.frame, group_by=base_t)).show()
+        status = view_t.add_column(agg=self.agg_fn(view_t.pos, view_t.frame, group_by=base_t))
         assert status.num_excs == 0
         _ = view_t.select(make_video(view_t.pos, view_t.agg)).group_by(base_t).show()
 
         # image cols computed with a window function currently need to be stored
         with pytest.raises(excs.Error):
-            view_t.add_column(agg2=agg_fn(view_t.pos, view_t.frame, group_by=base_t), stored=False)
+            view_t.add_column(agg2=self.agg_fn(view_t.pos, view_t.frame, group_by=base_t), stored=False)
 
         # reload from store
         cl = pxt.Client(reload=True)
         base_t, view_t = cl.get_table(base_t.get_name()), cl.get_table(view_t.get_name())
-        _ = view_t.select(agg_fn(view_t.pos, view_t.frame, group_by=base_t)).show()
+        _ = view_t.select(self.agg_fn(view_t.pos, view_t.frame, group_by=base_t)).show()

pixeltable/tests/test_view.py

@@ -1,3 +1,4 @@
+import datetime
 import logging
 
 import PIL
@@ -414,6 +415,10 @@ class TestView:
             v.order_by(v.c2).show(0),
             t.where(t.c2 < 10).order_by(t.c2).show(0))
 
+        # create views with filters containing date and datetime
+        _ = cl.create_view('test_view_2', t, filter=t.c5 >= datetime.date.today())
+        _ = cl.create_view('test_view_3', t, filter=t.c5 < datetime.datetime.now())
+
     def test_view_of_snapshot(self, test_client: pxt.Client) -> None:
         """Test view over a snapshot"""
         cl = test_client