pixeltable 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/tests/utils.py

@@ -1,16 +1,22 @@
 import datetime
 import glob
+import json
 import os
 from pathlib import Path
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 
 import numpy as np
 import pandas as pd
+import pytest
 
-import pixeltable as pt
+import pixeltable as pxt
+import pixeltable.type_system as ts
 from pixeltable import catalog
-from pixeltable.type_system import ColumnType, StringType, IntType, FloatType, BoolType, TimestampType
-from pixeltable.function import Function
+from pixeltable.dataframe import DataFrameResultSet
+from pixeltable.env import Env
+from pixeltable.type_system import \
+    ColumnType, StringType, IntType, FloatType, ArrayType, BoolType, TimestampType, JsonType, ImageType, VideoType
+
 
 def make_default_type(t: ColumnType.Type) -> ColumnType:
     if t == ColumnType.Type.STRING:
@@ -25,14 +31,19 @@ def make_default_type(t: ColumnType.Type) -> ColumnType:
         return TimestampType()
     assert False
 
-def make_tbl(db: pt.Db, name: str = 'test', col_names: List[str] = ['c1']) -> pt.MutableTable:
-    schema: List[catalog.Column] = []
+def make_tbl(cl: pxt.Client, name: str = 'test', col_names: Optional[List[str]] = None) -> catalog.InsertableTable:
+    if col_names is None:
+        col_names = ['c1']
+    schema: Dict[str, ts.ColumnType] = {}
     for i, col_name in enumerate(col_names):
-        schema.append(catalog.Column(f'{col_name}', make_default_type(ColumnType.Type(i % 5))))
-    return db.create_table(name, schema)
+        schema[f'{col_name}'] = make_default_type(ColumnType.Type(i % 5))
+    return cl.create_table(name, schema)
 
-def create_table_data(t: catalog.Table, col_names: List[str] = [], num_rows: int = 10) -> pd.DataFrame:
+def create_table_data(t: catalog.Table, col_names: Optional[List[str]] = None, num_rows: int = 10) -> List[Dict[str, Any]]:
+    if col_names is None:
+        col_names = []
     data: Dict[str, Any] = {}
+
     sample_dict = {
         'detections': [{
             'id': '637e8e073b28441a453564cf',
@@ -70,38 +81,138 @@ def create_table_data(t: catalog.Table, col_names: List[str] = [], num_rows: int
     }
 
     if len(col_names) == 0:
-        col_names = [c.name for c in t.columns]
+        col_names = [c.name for c in t.columns() if not c.is_computed]
 
+    col_types = t.column_types()
     for col_name in col_names:
-        col = t.cols_by_name[col_name]
+        col_type = col_types[col_name]
         col_data: Any = None
-        if col.col_type.is_string_type():
+        if col_type.is_string_type():
             col_data = ['test string'] * num_rows
-        if col.col_type.is_int_type():
-            col_data = np.random.randint(0, 100, size=num_rows)
-        if col.col_type.is_float_type():
-            col_data = np.random.random(size=num_rows) * 100
-        if col.col_type.is_bool_type():
+        if col_type.is_int_type():
+            col_data = np.random.randint(0, 100, size=num_rows).tolist()
+        if col_type.is_float_type():
+            col_data = (np.random.random(size=num_rows) * 100).tolist()
+        if col_type.is_bool_type():
             col_data = np.random.randint(0, 2, size=num_rows)
             col_data = [False if i == 0 else True for i in col_data]
-        if col.col_type.is_timestamp_type():
-            col_data = datetime.datetime.now()
-        if col.col_type.is_json_type():
+        if col_type.is_timestamp_type():
+            col_data = [datetime.datetime.now()] * num_rows
+        if col_type.is_json_type():
             col_data = [sample_dict] * num_rows
-        # TODO: implement this
-        assert not col.col_type.is_image_type()
-        assert not col.col_type.is_array_type()
-        data[col.name] = col_data
-    return pd.DataFrame(data=data)
+        if col_type.is_array_type():
+            col_data = [np.ones(col_type.shape, dtype=col_type.numpy_dtype()) for i in range(num_rows)]
+        if col_type.is_image_type():
+            image_path = get_image_files()[0]
+            col_data = [image_path for i in range(num_rows)]
+        if col_type.is_video_type():
+            video_path = get_video_files()[0]
+            col_data = [video_path for i in range(num_rows)]
+        data[col_name] = col_data
+    rows = [{col_name: data[col_name][i] for col_name in col_names} for i in range(num_rows)]
+    return rows
+
+def create_test_tbl(client: pxt.Client, name: str = 'test_tbl') -> catalog.Table:
+    schema = {
+        'c1': StringType(nullable=False),
+        'c1n': StringType(nullable=True),
+        'c2': IntType(nullable=False),
+        'c3': FloatType(nullable=False),
+        'c4': BoolType(nullable=False),
+        'c5': TimestampType(nullable=False),
+        'c6': JsonType(nullable=False),
+        'c7': JsonType(nullable=False),
+    }
+    t = client.create_table(name, schema, primary_key='c2')
+    t.add_column(c8=[[1, 2, 3], [4, 5, 6]])
+
+    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
+    rows = [
+        {
+            'c1': c1_data[i],
+            'c1n': c1_data[i] if i % 10 != 0 else None,
+            'c2': c2_data[i],
+            'c3': c3_data[i],
+            'c4': c4_data[i],
+            'c5': c5_data[i],
+            'c6': c6_data[i],
+            'c7': c7_data[i],
+        }
+        for i in range(num_rows)
+    ]
+    t.insert(rows)
+    return t
 
-def read_data_file(dir_name: str, file_name: str, path_col_names: List[str] = []) -> pd.DataFrame:
+def create_all_datatypes_tbl(test_client: pxt.Client) -> catalog.Table:
+    """ Creates a table with all supported datatypes.
+    """
+    schema = {
+        'row_id': IntType(nullable=False), # used for row selection
+        'c_array': ArrayType(shape=(10,),  dtype=FloatType(), nullable=True),
+        'c_bool': BoolType(nullable=True),
+        'c_float': FloatType(nullable=True),
+        'c_image': ImageType(nullable=True),
+        'c_int': IntType(nullable=True),
+        'c_json': JsonType(nullable=True),
+        'c_string': StringType(nullable=True),
+        'c_timestamp': TimestampType(nullable=True),
+        'c_video': VideoType(nullable=True),
+    }
+    tbl = test_client.create_table('all_datatype_tbl', schema)
+    example_rows = create_table_data(tbl, num_rows=11)
+
+    for i,r in enumerate(example_rows):
+        r['row_id'] = i # row_id
+
+    tbl.insert(example_rows)
+    return tbl
+
+def read_data_file(dir_name: str, file_name: str, path_col_names: Optional[List[str]] = None) -> List[Dict[str, Any]]:
     """
     Locate dir_name, create df out of file_name.
-    transform columns 'file_name' to column 'file_path' with absolute paths
     path_col_names: col names in csv file that contain file names; those will be converted to absolute paths
     by adding the path to 'file_name' as a prefix.
+    Returns:
+        tuple of (list of rows, list of column names)
     """
-    glob_result = glob.glob(f'{os.getcwd()}/**/{dir_name}', recursive=True)
+    if path_col_names is None:
+        path_col_names = []
+    tests_dir = os.path.dirname(__file__) # search with respect to tests/ dir
+    glob_result = glob.glob(f'{tests_dir}/**/{dir_name}', recursive=True)
     assert len(glob_result) == 1, f'Could not find {dir_name}'
     abs_path = Path(glob_result[0])
     data_file_path = abs_path / file_name
@@ -110,24 +221,54 @@ def read_data_file(dir_name: str, file_name: str, path_col_names: List[str] = []
     for col_name in path_col_names:
         assert col_name in df.columns
         df[col_name] = df.apply(lambda r: str(abs_path / r[col_name]), axis=1)
-        return df
+    return df.to_dict(orient='records')
+
+def get_video_files(include_bad_video=False) -> List[str]:
+    tests_dir = os.path.dirname(__file__) # search with respect to tests/ dir
+    glob_result = glob.glob(f'{tests_dir}/**/videos/*', recursive=True)
+    if not include_bad_video:
+        glob_result = [f for f in glob_result if 'bad_video' not in f]
+    return glob_result
 
-def get_video_files() -> List[str]:
-    glob_result = glob.glob(f'{os.getcwd()}/**/videos/*.mp4', recursive=True)
+def get_image_files() -> List[str]:
+    tests_dir = os.path.dirname(__file__) # search with respect to tests/ dir
+    glob_result = glob.glob(f'{tests_dir}/**/imagenette2-160/*', recursive=True)
     return glob_result
 
+def get_audio_files(include_bad_audio=False) -> List[str]:
+    tests_dir = os.path.dirname(__file__)
+    glob_result = glob.glob(f'{tests_dir}/**/audio/*', recursive=True)
+    if not include_bad_audio:
+        glob_result = [f for f in glob_result if 'bad_audio' not in f]
+    return glob_result
+
+def get_documents() -> List[str]:
+    tests_dir = os.path.dirname(__file__)
+    # for now, we can only handle .html and .md
+    return [p for p in glob.glob(f'{tests_dir}/**/documents/*', recursive=True) if not p.endswith('.pdf')]
+
+def get_sentences(n: int = 100) -> List[str]:
+    tests_dir = os.path.dirname(__file__)
+    path = glob.glob(f'{tests_dir}/**/jeopardy.json', recursive=True)[0]
+    with open(path, 'r') as f:
+        questions_list = json.load(f)
+    # this dataset contains \' around the questions
+    return [q['question'].replace("'", '') for q in questions_list[:n]]
+
+def assert_resultset_eq(r1: DataFrameResultSet, r2: DataFrameResultSet) -> None:
+    assert len(r1) == len(r2)
+    assert len(r1.column_names()) == len(r2.column_names())  # we don't care about the actual column names
+    r1_pd = r1.to_pandas()
+    r2_pd = r2.to_pandas()
+    for i in range(len(r1.column_names())):
+        # only compare column values
+        s1 = r1_pd.iloc[:, i]
+        s2 = r2_pd.iloc[:, i]
+        if s1.dtype == np.float64:
+            assert np.allclose(s1, s2)
+        else:
+            assert s1.equals(s2)
 
-class SumAggregator:
-    def __init__(self):
-        self.sum = 0
-    @classmethod
-    def make_aggregator(cls) -> 'SumAggregator':
-        return cls()
-    def update(self, val: int) -> None:
-        self.sum += val
-    def value(self) -> int:
-        return self.sum
-
-sum_uda = Function(
-    IntType(), [IntType()],
-    init_fn=SumAggregator.make_aggregator, update_fn=SumAggregator.update, value_fn=SumAggregator.value)
+def skip_test_if_not_installed(package) -> None:
+    if not Env.get().is_installed_package(package):
+        pytest.skip(f'Package `{package}` is not installed.')

pixeltable/tool/create_test_db_dump.py

@@ -0,0 +1,149 @@
+import datetime
+import json
+import logging
+import os
+import pathlib
+import subprocess
+
+import pgserver
+import toml
+
+import pixeltable as pxt
+import pixeltable.metadata as metadata
+from pixeltable.env import Env
+from pixeltable.type_system import \
+    StringType, IntType, FloatType, BoolType, TimestampType, JsonType
+
+_logger = logging.getLogger('pixeltable')
+
+
+class Dumper:
+
+    def __init__(self, output_dir='target', db_name='pxtdump') -> None:
+        self.output_dir = pathlib.Path(output_dir)
+        shared_home = pathlib.Path(os.environ.get('PIXELTABLE_HOME', '~/.pixeltable')).expanduser()
+        mock_home_dir = self.output_dir / '.pixeltable'
+        mock_home_dir.mkdir(parents=True, exist_ok=True)
+        os.environ['PIXELTABLE_HOME'] = str(mock_home_dir)
+        os.environ['PIXELTABLE_CONFIG'] = str(shared_home / 'config.yaml')
+        os.environ['PIXELTABLE_DB'] = db_name
+        os.environ['PIXELTABLE_PGDATA'] = str(shared_home / 'pgdata')
+
+        Env.get().set_up(reinit_db=True)
+        self.cl = pxt.Client()
+        self.cl.logging(level=logging.DEBUG, to_stdout=True)
+
+    def dump_db(self) -> None:
+        md_version = metadata.VERSION
+        dump_file = self.output_dir / f'pixeltable-v{md_version:03d}-test.dump.gz'
+        _logger.info(f'Creating database dump at: {dump_file}')
+        pg_package_dir = os.path.dirname(pgserver.__file__)
+        pg_dump_binary = f'{pg_package_dir}/pginstall/bin/pg_dump'
+        _logger.info(f'Using pg_dump binary at: {pg_dump_binary}')
+        with open(dump_file, 'wb') as dump:
+            pg_dump_process = subprocess.Popen(
+                [pg_dump_binary, Env.get().db_url, '-U', 'postgres', '-Fc'],
+                stdout=subprocess.PIPE
+            )
+            subprocess.run(
+                ["gzip", "-9"],
+                stdin=pg_dump_process.stdout,
+                stdout=dump,
+                check=True
+            )
+        info_file = self.output_dir / f'pixeltable-v{md_version:03d}-test-info.toml'
+        git_sha = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('ascii').strip()
+        user = os.environ.get('USER', os.environ.get('USERNAME'))
+        info_dict = {'pixeltable-dump': {
+            'metadata-version': md_version,
+            'git-sha': git_sha,
+            'datetime': datetime.datetime.utcnow(),
+            'user': user
+        }}
+        with open(info_file, 'w') as info:
+            toml.dump(info_dict, info)
+
+    # TODO: Add additional features to the test DB dump (ideally it should exercise
+    # every major pixeltable DB feature)
+    def create_tables(self) -> None:
+        schema = {
+            'c1': StringType(nullable=False),
+            'c1n': StringType(nullable=True),
+            'c2': IntType(nullable=False),
+            'c3': FloatType(nullable=False),
+            'c4': BoolType(nullable=False),
+            'c5': TimestampType(nullable=False),
+            'c6': JsonType(nullable=False),
+            'c7': JsonType(nullable=False),
+        }
+        t = self.cl.create_table('sample_table', schema, primary_key='c2')
+        t.add_column(c8=[[1, 2, 3], [4, 5, 6]])
+
+        # Add columns for .astype converters to ensure they're persisted properly
+        t.add_column(c2_as_float=t.c2.astype(FloatType()))
+
+        # Add columns for .apply
+        t.add_column(c2_to_string=t.c2.apply(str))
+        t.add_column(c6_to_string=t.c6.apply(json.dumps))
+        t.add_column(c6_back_to_json=t.c6_to_string.apply(json.loads))
+
+        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
+        rows = [
+            {
+                'c1': c1_data[i],
+                'c1n': c1_data[i] if i % 10 != 0 else None,
+                'c2': c2_data[i],
+                'c3': c3_data[i],
+                'c4': c4_data[i],
+                'c5': c5_data[i],
+                'c6': c6_data[i],
+                'c7': c7_data[i],
+            }
+            for i in range(num_rows)
+        ]
+        t.insert(rows)
+
+
+def main() -> None:
+    _logger.info("Creating pixeltable test artifact.")
+    dumper = Dumper()
+    dumper.create_tables()
+    dumper.dump_db()
+
+
+if __name__ == "__main__":
+    main()