pixeltable 0.0.0__py3-none-any.whl

pixeltable/store.py

@@ -0,0 +1,424 @@
+from __future__ import annotations
+
+import abc
+import logging
+import os
+import sys
+import urllib.parse
+import urllib.request
+import warnings
+from typing import Optional, Dict, Any, List, Tuple, Set
+
+import sqlalchemy as sql
+from tqdm import tqdm, TqdmWarning
+
+import pixeltable.catalog as catalog
+import pixeltable.env as env
+from pixeltable import exprs
+import pixeltable.exceptions as excs
+from pixeltable.exec import ExecNode
+from pixeltable.metadata import schema
+from pixeltable.type_system import StringType
+from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils.sql import log_stmt, log_explain
+
+_logger = logging.getLogger('pixeltable')
+
+
+class StoreBase:
+    """Base class for stored tables
+
+    Each row has the following system columns:
+    - rowid columns: one or more columns that identify a user-visible row across all versions
+    - v_min: version at which the row was created
+    - v_max: version at which the row was deleted (or MAX_VERSION if it's still live)
+    """
+
+    def __init__(self, tbl_version: catalog.TableVersion):
+        self.tbl_version = tbl_version
+        self.sa_md = sql.MetaData()
+        self.sa_tbl: Optional[sql.Table] = None
+        self.create_sa_tbl()
+
+    def pk_columns(self) -> List[sql.Column]:
+        return self._pk_columns
+
+    def rowid_columns(self) -> List[sql.Column]:
+        return self._pk_columns[:-1]
+
+    @abc.abstractmethod
+    def _create_rowid_columns(self) -> List[sql.Column]:
+        """Create and return rowid columns"""
+        pass
+
+    @abc.abstractmethod
+    def _create_system_columns(self) -> List[sql.Column]:
+        """Create and return system columns"""
+        rowid_cols = self._create_rowid_columns()
+        self.v_min_col = sql.Column('v_min', sql.BigInteger, nullable=False)
+        self.v_max_col = \
+            sql.Column('v_max', sql.BigInteger, nullable=False, server_default=str(schema.Table.MAX_VERSION))
+        self._pk_columns = [*rowid_cols, self.v_min_col]
+        return [*rowid_cols, self.v_min_col, self.v_max_col]
+
+
+    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()
+        idxs: List[sql.Index] = []
+        for col in [c for c in self.tbl_version.cols if c.is_stored]:
+            # re-create sql.Column for each column, regardless of whether it already has sa_col set: it was bound
+            # to the last sql.Table version we created and cannot be reused
+            col.create_sa_cols()
+            all_cols.append(col.sa_col)
+            if col.records_errors:
+                all_cols.append(col.sa_errormsg_col)
+                all_cols.append(col.sa_errortype_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
+            #   filter on)
+            if (col.col_type.is_scalar_type() and not col.col_type.is_string_type()) \
+                    or (col.col_type.is_media_type() and not col.is_computed):
+                # index names need to be unique within the Postgres instance
+                idx_name = f'idx_{col.id}_{self.tbl_version.id.hex}'
+                idxs.append(sql.Index(idx_name, col.sa_col))
+
+        if self.sa_tbl is not None:
+            # if we're called in response to a schema change, we need to remove the old table first
+            self.sa_md.remove(self.sa_tbl)
+
+        # index for all system columns:
+        # - base x view joins can be executed as merge joins
+        # - speeds up ORDER BY rowid DESC
+        # - allows filtering for a particular table version in index scan
+        idx_name = f'sys_cols_idx_{self.tbl_version.id.hex}'
+        idxs.append(sql.Index(idx_name, *system_cols))
+        # v_min/v_max indices: speeds up base table scans needed to propagate a base table insert or delete
+        idx_name = f'vmin_idx_{self.tbl_version.id.hex}'
+        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using='brin'))
+        idx_name = f'vmax_idx_{self.tbl_version.id.hex}'
+        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using='brin'))
+
+        self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
+
+    @abc.abstractmethod
+    def _rowid_join_predicate(self) -> sql.ClauseElement:
+        """Return predicate for rowid joins to all bases"""
+        pass
+
+    @abc.abstractmethod
+    def _storage_name(self) -> str:
+        """Return the name of the data store table"""
+        pass
+
+    def _move_tmp_media_file(self, file_url: Optional[str], col: catalog.Column, v_min: int) -> str:
+        """Move tmp media file with given url to Env.media_dir and return new url, or given url if not a tmp_dir file"""
+        pxt_tmp_dir = str(env.Env.get().tmp_dir)
+        if file_url is None:
+            return None
+        parsed = urllib.parse.urlparse(file_url)
+        # We should never be passed a local file path here. The "len > 1" ensures that Windows
+        # file paths aren't mistaken for URLs with a single-character scheme.
+        assert len(parsed.scheme) > 1
+        if parsed.scheme != 'file':
+            # remote url
+            return file_url
+        file_path = urllib.parse.unquote(urllib.request.url2pathname(parsed.path))
+        if not file_path.startswith(pxt_tmp_dir):
+            # not a tmp file
+            return file_url
+        _, ext = os.path.splitext(file_path)
+        new_path = str(MediaStore.prepare_media_path(self.tbl_version.id, col.id, v_min, ext=ext))
+        os.rename(file_path, new_path)
+        new_file_url = urllib.parse.urljoin('file:', urllib.request.pathname2url(new_path))
+        return new_file_url
+
+    def _move_tmp_media_files(
+            self, table_rows: List[Dict[str, Any]], media_cols: List[catalog.Column], v_min: int
+    ) -> None:
+        """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.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],
+            exc_col_ids: Set[int], v_min: int
+    ) -> Tuple[Dict[str, Any], int]:
+        """Return Tuple[complete table row, # of exceptions] for insert()
+        Creates a row that includes the PK columns, with the values from input_row.pk.
+        Returns:
+            Tuple[complete table row, # of exceptions]
+        """
+        table_row, num_excs = row_builder.create_table_row(input_row, exc_col_ids)
+
+        assert input_row.pk is not None and len(input_row.pk) == len(self._pk_columns)
+        for pk_col, pk_val in zip(self._pk_columns, input_row.pk):
+            if pk_col == self.v_min_col:
+                table_row[pk_col.name] = v_min
+            else:
+                table_row[pk_col.name] = pk_val
+
+        return table_row, num_excs
+
+    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)
+        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
+
+    def create(self, conn: sql.engine.Connection) -> None:
+        self.sa_md.create_all(bind=conn)
+
+    def drop(self, conn: sql.engine.Connection) -> None:
+        """Drop store table"""
+        self.sa_md.drop_all(bind=conn)
+
+    def add_column(self, col: catalog.Column, conn: sql.engine.Connection) -> None:
+        """Add column(s) to the store-resident table based on a catalog column
+
+        Note that a computed catalog column will require two extra columns (for the computed value and for the error
+        message).
+        """
+        assert col.is_stored
+        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.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_store_name()} VARCHAR DEFAULT NULL')
+            conn.execute(sql.text(stmt))
+            stmt = (f'ALTER TABLE {self._storage_name()} '
+                    f'ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL')
+            conn.execute(sql.text(stmt))
+            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: 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))
+
+    def load_column(
+            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
+
+        Returns:
+            number of rows with exceptions
+        Raises:
+            sql.exc.DBAPIError if there was an error during SQL execution
+        """
+        num_excs = 0
+        num_rows = 0
+        for row_batch in exec_plan:
+            num_rows += len(row_batch)
+            for result_row in row_batch:
+                values_dict: Dict[sql.Column, Any] = {}
+
+                if col.is_computed:
+                    if result_row.has_exc(value_expr_slot_idx):
+                        num_excs += 1
+                        value_exc = result_row.get_exc(value_expr_slot_idx)
+                        # we store a NULL value and record the exception/exc type
+                        error_type = type(value_exc).__name__
+                        error_msg = str(value_exc)
+                        values_dict = {
+                            col.sa_col: None,
+                            col.sa_errortype_col: error_type,
+                            col.sa_errormsg_col: error_msg
+                        }
+                    else:
+                        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}
+
+                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)
+                log_stmt(_logger, update_stmt)
+                conn.execute(update_stmt)
+
+        return num_excs
+
+    def insert_rows(
+            self, exec_plan: ExecNode, conn: sql.engine.Connection, v_min: Optional[int] = None
+    ) -> Tuple[int, int, Set[int]]:
+        """Insert rows into the store table and update the catalog table's md
+        Returns:
+            number of inserted rows, number of exceptions, set of column ids that have exceptions
+        """
+        assert v_min is not None
+        exec_plan.ctx.conn = conn
+        batch_size = 16  # TODO: is this a good batch size?
+        # TODO: total?
+        num_excs = 0
+        num_rows = 0
+        cols_with_excs: Set[int] = set()
+        progress_bar: Optional[tqdm] = None  # create this only after we started executing
+        row_builder = exec_plan.row_builder
+        media_cols = [info.col for info in row_builder.table_columns if info.col.col_type.is_media_type()]
+        try:
+            exec_plan.open()
+            for row_batch in exec_plan:
+                num_rows += len(row_batch)
+                for batch_start_idx in range(0, len(row_batch), batch_size):
+                    # compute batch of rows and convert them into table rows
+                    table_rows: List[Dict[str, Any]] = []
+                    for row_idx in range(batch_start_idx, min(batch_start_idx + batch_size, len(row_batch))):
+                        row = row_batch[row_idx]
+                        table_row, num_row_exc = \
+                            self._create_table_row(row, row_builder, media_cols, cols_with_excs, v_min=v_min)
+                        num_excs += num_row_exc
+                        table_rows.append(table_row)
+                        if progress_bar is None:
+                            warnings.simplefilter("ignore", category=TqdmWarning)
+                            progress_bar = tqdm(
+                                desc=f'Inserting rows into `{self.tbl_version.name}`',
+                                unit=' rows',
+                                ncols=100,
+                                file=sys.stdout
+                            )
+                        progress_bar.update(1)
+                    self._move_tmp_media_files(table_rows, media_cols, v_min)
+                    conn.execute(sql.insert(self.sa_tbl), table_rows)
+            if progress_bar is not None:
+                progress_bar.close()
+            return num_rows, num_excs, cols_with_excs
+        finally:
+            exec_plan.close()
+
+    def _versions_clause(self, versions: List[Optional[int]], match_on_vmin: bool) -> sql.ClauseElement:
+        """Return filter for base versions"""
+        v = versions[0]
+        if v is None:
+            # we're looking at live rows
+            clause = sql.and_(self.v_min_col <= self.tbl_version.version, self.v_max_col == schema.Table.MAX_VERSION)
+        else:
+            # we're looking at a specific version
+            clause = self.v_min_col == v if match_on_vmin else self.v_max_col == v
+        if len(versions) == 1:
+            return clause
+        return sql.and_(clause, self.base._versions_clause(versions[1:], match_on_vmin))
+
+    def delete_rows(
+            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
+                current_version parameter)
+            match_on_vmin: if True, match exact versions on v_min; if False, match on v_max
+            where_clause: if not None, also apply where_clause
+        Returns:
+            number of deleted rows
+        """
+        where_clause = sql.true() if where_clause is None else where_clause
+        where_clause = sql.and_(
+            self.v_min_col < current_version,
+            self.v_max_col == schema.Table.MAX_VERSION,
+            where_clause)
+        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(set_clause) \
+            .where(where_clause) \
+            .where(rowid_join_clause) \
+            .where(base_versions_clause)
+        log_explain(_logger, stmt, conn)
+        status = conn.execute(stmt)
+        return status.rowcount
+
+
+class StoreTable(StoreBase):
+    def __init__(self, tbl_version: catalog.TableVersion):
+        assert not tbl_version.is_view()
+        super().__init__(tbl_version)
+
+    def _create_rowid_columns(self) -> List[sql.Column]:
+        self.rowid_col = sql.Column('rowid', sql.BigInteger, nullable=False)
+        return [self.rowid_col]
+
+    def _storage_name(self) -> str:
+        return f'tbl_{self.tbl_version.id.hex}'
+
+    def _rowid_join_predicate(self) -> sql.ClauseElement:
+        return sql.true()
+
+
+class StoreView(StoreBase):
+    def __init__(self, catalog_view: catalog.TableVersion):
+        assert catalog_view.is_view()
+        self.base = catalog_view.base.store_tbl
+        super().__init__(catalog_view)
+
+    def _create_rowid_columns(self) -> List[sql.Column]:
+        # a view row corresponds directly to a single base row, which means it needs to duplicate its rowid columns
+        self.rowid_cols = [sql.Column(c.name, c.type) for c in self.base.rowid_columns()]
+        return self.rowid_cols
+
+    def _storage_name(self) -> str:
+        return f'view_{self.tbl_version.id.hex}'
+
+    def _rowid_join_predicate(self) -> sql.ClauseElement:
+        return sql.and_(
+            self.base._rowid_join_predicate(),
+            *[c1 == c2 for c1, c2 in zip(self.rowid_columns(), self.base.rowid_columns())])
+
+class StoreComponentView(StoreView):
+    """A view that stores components of its base, as produced by a ComponentIterator
+
+    PK: now also includes pos, the position returned by the ComponentIterator for the base row identified by base_rowid
+    """
+    def __init__(self, catalog_view: catalog.TableVersion):
+        super().__init__(catalog_view)
+
+    def _create_rowid_columns(self) -> List[sql.Column]:
+        # each base row is expanded into n view rows
+        self.rowid_cols = [sql.Column(c.name, c.type) for c in self.base.rowid_columns()]
+        # name of pos column: avoid collisions with bases' pos columns
+        self.pos_col = sql.Column(f'pos_{len(self.rowid_cols) - 1}', sql.BigInteger, nullable=False)
+        self.pos_col_idx = len(self.rowid_cols)
+        self.rowid_cols.append(self.pos_col)
+        return self.rowid_cols
+
+    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
+
+    def _rowid_join_predicate(self) -> sql.ClauseElement:
+        return sql.and_(
+            self.base._rowid_join_predicate(),
+            *[c1 == c2 for c1, c2 in zip(self.rowid_columns()[:-1], self.base.rowid_columns())])

pixeltable/tool/create_test_db_dump.py

@@ -0,0 +1,184 @@
+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.func import Batch
+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().configure_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 = pxt.create_table('sample_table', schema, primary_key='c2')
+
+        # Add columns for InlineArray and InlineDict
+        t.add_column(c8=[[1, 2, 3], [4, 5, 6]])
+        t.add_column(c9=[['a', 'b', 'c'], ['d', 'e', 'f']])
+        t.add_column(c10=[t.c1, [t.c1n, t.c2]])
+        t.add_column(c11={'int': 22, 'dict': {'key': 'val'}, 'expr': t.c1})
+
+        # InPredicate
+        t.add_column(isin_1=t.c1.isin(['test string 1', 'test string 2', 'test string 3']))
+        t.add_column(isin_2=t.c2.isin([1, 2, 3, 4, 5]))
+        t.add_column(isin_3=t.c2.isin(t.c6.f5))
+
+        # 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)
+        pxt.create_dir('views')
+        v = pxt.create_view('views.sample_view', t, filter=(t.c2 < 50))
+        _ = pxt.create_view('views.sample_snapshot', t, filter=(t.c2 >= 75), is_snapshot=True)
+        e = pxt.create_view('views.empty_view', t, filter=t.c2 == 4171780)
+        assert e.count() == 0
+        # Computed column using a library function
+        v['str_format'] = pxt.functions.string.str_format('{0} {key}', t.c1, key=t.c1)
+        # Computed column using a bespoke stored udf
+        v['test_udf'] = test_udf_stored(t.c2)
+        # Computed column using a batched function
+        # (apply this to the empty view, since it's a "heavyweight" function)
+        e['batched'] = pxt.functions.huggingface.clip_text(t.c1, model_id='openai/clip-vit-base-patch32')
+        # computed column using a stored batched function
+        v['test_udf_batched'] = test_udf_stored_batched(t.c1, upper=False)
+        # astype
+        v['astype'] = t.c1.astype(pxt.FloatType())
+
+
+@pxt.udf(_force_stored=True)
+def test_udf_stored(n: int) -> int:
+    return n + 1
+
+
+@pxt.udf(batch_size=4, _force_stored=True)
+def test_udf_stored_batched(strings: Batch[str], *, upper: bool = True) -> Batch[str]:
+    return [string.upper() if upper else string.lower() for string in strings]
+
+
+def main() -> None:
+    _logger.info("Creating pixeltable test artifact.")
+    dumper = Dumper()
+    dumper.create_tables()
+    dumper.dump_db()
+
+
+if __name__ == "__main__":
+    main()

pixeltable/tool/create_test_video.py

@@ -0,0 +1,81 @@
+import av
+import PIL.Image
+import PIL.ImageDraw
+import PIL.ImageFont
+
+from pathlib import Path
+from typing import Optional
+import tempfile
+import math
+
+def create_test_video(
+    frame_count: int,
+    frame_rate: float = 1.0,
+    frame_width: int = 224,
+    aspect_ratio: str = '16:9',
+    frame_height: Optional[int] = None,
+    output_path: Optional[Path] = None,
+    font_file: str = '/Library/Fonts/Arial Unicode.ttf',
+) -> Path:
+    """
+        Creates an .mp4 video file such as the ones in /tests/data/test_videos
+        The video contains a frame number in each frame (for visual sanity check).
+
+        Args:
+            frame_count: Number of frames to create
+            frame_rate: Frame rate of the video
+            frame_width (int): Width in pixels of the video frame. Note: cost of decoding increases dramatically
+                with frame width * frame height.
+            aspect_ratio: Aspect ratio (width/height) of the video frames string of form 'width:height'
+            frame_height: Height of the video frame, if given, aspect_ratio is ignored
+            output_path: Path to save the video file
+            font_file: Path to the font file used for text.
+    """
+
+    if output_path is None:
+        output_path = Path(tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name)
+
+    parts = [int(p) for p in aspect_ratio.split(':')]
+    assert len(parts) == 2
+    aspect_ratio = parts[0] / parts[1]
+
+    if frame_height is None:
+        frame_height = math.ceil(frame_width / aspect_ratio)
+
+    frame_size = (frame_width, frame_height)
+
+    font_size = min(frame_height, frame_width) // 4
+    font = PIL.ImageFont.truetype(font=font_file, size=font_size)
+    font_fill = 0xFFFFFF  # white
+    frame_color = 0xFFFFFF - font_fill  # black
+    # Create a video container
+    container = av.open(str(output_path), mode='w')
+
+    # Add a video stream
+    stream = container.add_stream('h264', rate=frame_rate)
+    stream.width, stream.height = frame_size
+    stream.pix_fmt = 'yuv420p'
+
+    for frame_number in range(frame_count):
+        # Create an image with a number in it
+        image = PIL.Image.new('RGB', frame_size, color=frame_color)
+        draw = PIL.ImageDraw.Draw(image)
+        # Optionally, add a font here if you have one
+        text = str(frame_number)
+        _, _, text_width, text_height = draw.textbbox((0, 0), text, font=font)
+        text_position = ((frame_size[0] - text_width) // 2, (frame_size[1] - text_height) // 2)
+        draw.text(text_position, text, font=font, fill=font_fill)
+
+        # Convert the PIL image to an AVFrame
+        frame = av.VideoFrame.from_image(image)
+
+        # Encode and write the frame
+        for packet in stream.encode(frame):
+            container.mux(packet)
+
+    # Flush and close the stream
+    for packet in stream.encode():
+        container.mux(packet)
+
+    container.close()
+    return output_path