pixeltable 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

pixeltable/share/packager.py

@@ -0,0 +1,218 @@
+import io
+import json
+import logging
+import tarfile
+import urllib.parse
+import urllib.request
+import uuid
+from pathlib import Path
+from typing import Any, Iterator
+
+import more_itertools
+import numpy as np
+import pyarrow as pa
+import pyiceberg.catalog
+
+import pixeltable as pxt
+import pixeltable.type_system as ts
+from pixeltable import exprs
+from pixeltable.env import Env
+from pixeltable.utils.arrow import PXT_TO_PA_TYPES
+from pixeltable.utils.iceberg import sqlite_catalog
+
+_logger = logging.getLogger('pixeltable')
+
+
+class TablePackager:
+    """
+    Packages a pixeltable Table into a tarball containing Iceberg tables and media files. The structure of the tarball
+    is as follows:
+
+    warehouse/catalog.db  # sqlite Iceberg catalog
+    warehouse/pxt.db/**  # Iceberg metadata and data files (parquet/avro/json)
+    media/**  # Local media files
+
+    If the table being archived is a view, then the Iceberg catalog will contain separate tables for the view and each
+    of its ancestors. All rows will be exported with additional _rowid and _v_min columns. Currently, only the most
+    recent version of the table can be exported, and only the full table contents.
+
+    If the table contains media columns, they are handled as follows:
+    - If a media file has an external URL (any URL scheme other than file://), then the URL will be preserved as-is and
+      stored in the Iceberg table.
+    - If a media file is a local file, then it will be copied into the tarball as a file of the form
+      'media/{uuid}{extension}', and the Iceberg table will contain the ephemeral URI 'pxtmedia://{uuid}{extension}'.
+    """
+
+    table: pxt.Table  # The table to be packaged
+    tmp_dir: Path  # Temporary directory where the package will reside
+    iceberg_catalog: pyiceberg.catalog.Catalog
+    media_files: dict[Path, str]  # Mapping from local media file paths to their tarball names
+
+    def __init__(self, table: pxt.Table) -> None:
+        self.table = table
+        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.media_files = {}
+
+    def package(self) -> Path:
+        """
+        Export the table to a tarball containing Iceberg tables and media files.
+        """
+        assert not self.tmp_dir.exists()  # Packaging can only be done once per TablePackager instance
+        _logger.info(f"Packaging table '{self.table._path}' and its ancestors in: {self.tmp_dir}")
+        self.tmp_dir.mkdir()
+        self.iceberg_catalog = sqlite_catalog(self.tmp_dir / 'warehouse')
+        ancestors = [self.table] + self.table._bases
+        for t in ancestors:
+            _logger.info(f"Exporting table '{t._path}'.")
+            self.__export_table(t)
+        _logger.info(f'Building archive.')
+        bundle_path = self.__build_tarball()
+        _logger.info(f'Packaging complete: {bundle_path}')
+        return bundle_path
+
+    def __export_table(self, t: pxt.Table) -> None:
+        """
+        Exports the data from `t` into an Iceberg table.
+        """
+        # First generate a select list for the data we want to extract from `t`. This includes:
+        # - all stored columns, including computed columns;
+        # - errortype and errormsg fields whenever they're defined.
+        # We select only those columns that are defined in this table (columns inherited from ancestor tables will be
+        # handled separately).
+        # For media columns, we substitute `col.fileurl` so that we always get the URL (which may be a file:// URL;
+        # these will be specially handled later)
+        select_exprs: dict[str, exprs.Expr] = {}
+
+        # As we generate the select list, we construct a separate list of column types. We can't rely on df._schema
+        # to get the column types, since we'll be substituting `fileurl`s for media columns.
+        actual_col_types: list[ts.ColumnType] = []
+
+        for col_name, col in t._tbl_version.cols_by_name.items():
+            if not col.is_stored:
+                continue
+            if col.col_type.is_media_type():
+                select_exprs[col_name] = t[col_name].fileurl
+            else:
+                select_exprs[col_name] = t[col_name]
+            actual_col_types.append(col.col_type)
+            if col.records_errors:
+                select_exprs[f'{col_name}_errortype'] = t[col_name].errortype
+                actual_col_types.append(ts.StringType())
+                select_exprs[f'{col_name}_errormsg'] = t[col_name].errormsg
+                actual_col_types.append(ts.StringType())
+
+        # Run the select() on `self.table`, not `t`, so that we export only those rows that are actually present in
+        # `self.table`.
+        df = self.table.select(**select_exprs)
+        namespace = self.__iceberg_namespace(t)
+        self.iceberg_catalog.create_namespace_if_not_exists(namespace)
+        iceberg_schema = self.__to_iceberg_schema(df._schema)
+        iceberg_tbl = self.iceberg_catalog.create_table(f'{namespace}.{t._name}', schema=iceberg_schema)
+
+        # Populate the Iceberg table with data.
+        # The data is first loaded from the DataFrame into a sequence of pyarrow tables, batched in order to avoid
+        # excessive memory usage. The pyarrow tables are then amalgamated into the (single) Iceberg table on disk.
+        for pa_table in self.__to_pa_tables(df, actual_col_types, iceberg_schema):
+            iceberg_tbl.append(pa_table)
+
+    @classmethod
+    def __iceberg_namespace(cls, table: pxt.Table) -> str:
+        """
+        Iceberg tables must have a namespace, which cannot be the empty string, so we prepend `pxt` to the table path.
+        """
+        parent_path = table._parent._path
+        if len(parent_path) == 0:
+            return 'pxt'
+        else:
+            return f'pxt.{parent_path}'
+
+    # The following methods are responsible for schema and data conversion from Pixeltable to Iceberg. Some of this
+    # logic might be consolidated into arrow.py and unified with general Parquet export, but there are several
+    # major differences:
+    # - Iceberg has no array type; we export all arrays as binary blobs
+    # - We include _rowid and _v_min columns in the Iceberg table
+    # - Media columns are handled specially as indicated above
+
+    @classmethod
+    def __to_iceberg_schema(cls, pxt_schema: dict[str, ts.ColumnType]) -> pa.Schema:
+        entries = [(name, cls.__to_iceberg_type(col_type)) for name, col_type in pxt_schema.items()]
+        entries.append(('_rowid', pa.list_(pa.int64())))
+        entries.append(('_v_min', pa.int64()))
+        return pa.schema(entries)  # type: ignore[arg-type]
+
+    @classmethod
+    def __to_iceberg_type(cls, col_type: ts.ColumnType) -> pa.DataType:
+        if col_type.is_array_type():
+            return pa.binary()
+        if col_type.is_media_type():
+            return pa.string()
+        return PXT_TO_PA_TYPES.get(col_type.__class__)
+
+    def __to_pa_tables(
+        self,
+        df: pxt.DataFrame,
+        actual_col_types: list[pxt.ColumnType],
+        arrow_schema: pa.Schema,
+        batch_size: int = 1_000,
+    ) -> Iterator[pa.Table]:
+        """
+        Load a DataFrame as a sequence of pyarrow tables. The pyarrow tables are batched into smaller chunks
+        to avoid excessive memory usage.
+        """
+        for rows in more_itertools.batched(self.__to_pa_rows(df, actual_col_types), batch_size):
+            cols = {col_name: [row[idx] for row in rows] for idx, col_name in enumerate(df._schema.keys())}
+            cols['_rowid'] = [row[-2] for row in rows]
+            cols['_v_min'] = [row[-1] for row in rows]
+            yield pa.Table.from_pydict(cols, schema=arrow_schema)
+
+    def __to_pa_rows(self, df: pxt.DataFrame, actual_col_types: list[pxt.ColumnType]) -> Iterator[list]:
+        for row in df._exec():
+            vals = [row[e.slot_idx] for e in df._select_list_exprs]
+            result = [self.__to_pa_value(val, col_type) for val, col_type in zip(vals, actual_col_types)]
+            result.append(row.rowid)
+            result.append(row.v_min)
+            yield result
+
+    def __to_pa_value(self, val: Any, col_type: ts.ColumnType) -> Any:
+        if val is None:
+            return None
+        if col_type.is_array_type():
+            # Export arrays as binary
+            assert isinstance(val, np.ndarray)
+            arr = io.BytesIO()
+            np.save(arr, val)
+            return arr.getvalue()
+        if col_type.is_json_type():
+            # Export JSON as strings
+            return json.dumps(val)
+        if col_type.is_media_type():
+            # Handle media files as described above
+            assert isinstance(val, str)  # Media columns are always referenced by `fileurl`
+            return self.__process_media_url(val)
+        return val
+
+    def __process_media_url(self, url: str) -> str:
+        parsed_url = urllib.parse.urlparse(url)
+        if parsed_url.scheme == 'file':
+            # It's the URL of a local file. Replace it with a pxtmedia:// URI.
+            # (We can't use an actual pxt:// URI, because the eventual pxt:// table name might not be known at this
+            # time. The pxtmedia:// URI serves as a relative reference into the tarball that can be replaced with an
+            # actual URL when the table is reconstituted.)
+            path = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed_url.path)))
+            if path not in self.media_files:
+                # Create a new entry in the `media_files` dict so that we can copy the file into the tarball later.
+                dest_name = f'{uuid.uuid4().hex}{path.suffix}'
+                self.media_files[path] = dest_name
+            return f'pxtmedia://{self.media_files[path]}'
+        # For any type of URL other than a local file, just return the URL as-is.
+        return url
+
+    def __build_tarball(self) -> Path:
+        bundle_path = self.tmp_dir / 'bundle.tar.bz2'
+        with tarfile.open(bundle_path, 'w:bz2') as tf:
+            # Add the Iceberg warehouse dir (including the catalog)
+            tf.add(self.tmp_dir / 'warehouse', arcname='warehouse', recursive=True)
+            # Add the media files
+            for src_file, dest_name in self.media_files.items():
+                tf.add(src_file, arcname=f'media/{dest_name}')
+        return bundle_path

pixeltable/store.py

@@ -32,6 +32,7 @@ class StoreBase:
     - 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)
     """
+
     tbl_version: catalog.TableVersion
     sa_md: sql.MetaData
     sa_tbl: Optional[sql.Table]
@@ -65,8 +66,9 @@ class StoreBase:
         """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.v_max_col = sql.Column(
+            'v_max', sql.BigInteger, nullable=False, server_default=str(schema.Table.MAX_VERSION)
+        )
         self._pk_cols = [*rowid_cols, self.v_min_col]
         return [*rowid_cols, self.v_min_col, self.v_max_col]
 
@@ -134,7 +136,7 @@ class StoreBase:
         return new_file_url
 
     def _move_tmp_media_files(
-            self, table_rows: list[dict[str, Any]], media_cols: list[catalog.Column], v_min: int
+        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:
@@ -143,7 +145,7 @@ class StoreBase:
                 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, exc_col_ids: set[int], pk: tuple[int, ...]
+        self, input_row: exprs.DataRow, row_builder: exprs.RowBuilder, exc_col_ids: set[int], pk: tuple[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.
@@ -193,11 +195,13 @@ class StoreBase:
         added_storage_cols = [col.store_name()]
         if col.records_errors:
             # we also need to create the errormsg and errortype storage cols
-            stmt = sql.text(f'ALTER TABLE {self._storage_name()} '
-                            f'ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL')
+            stmt = sql.text(
+                f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL'
+            )
             conn.execute(stmt)
-            stmt = sql.text(f'ALTER TABLE {self._storage_name()} '
-                            f'ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL')
+            stmt = sql.text(
+                f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL'
+            )
             conn.execute(stmt)
             added_storage_cols.extend([col.errormsg_store_name(), col.errortype_store_name()])
         self.create_sa_tbl()
@@ -219,7 +223,7 @@ class StoreBase:
         exec_plan: ExecNode,
         value_expr_slot_idx: int,
         conn: sql.engine.Connection,
-        on_error: Literal['abort', 'ignore']
+        on_error: Literal['abort', 'ignore'],
     ) -> int:
         """Update store column of a computed column with values produced by an execution plan
 
@@ -295,10 +299,9 @@ class StoreBase:
                 update_stmt = update_stmt.where(pk_col == tmp_pk_col)
             update_stmt = update_stmt.values({col.sa_col: tmp_val_col})
             if col.records_errors:
-                update_stmt = update_stmt.values({
-                    col.sa_errortype_col: tmp_errortype_col,
-                    col.sa_errormsg_col: tmp_errormsg_col
-                })
+                update_stmt = update_stmt.values(
+                    {col.sa_errortype_col: tmp_errortype_col, col.sa_errormsg_col: tmp_errormsg_col}
+                )
             log_explain(_logger, update_stmt, conn)
             conn.execute(update_stmt)
 
@@ -308,8 +311,13 @@ class StoreBase:
         return num_excs
 
     def insert_rows(
-            self, exec_plan: ExecNode, conn: sql.engine.Connection, v_min: Optional[int] = None,
-            show_progress: bool = True, rowids: Optional[Iterator[int]] = None, abort_on_exc: bool = False
+        self,
+        exec_plan: ExecNode,
+        conn: sql.engine.Connection,
+        v_min: Optional[int] = None,
+        show_progress: bool = True,
+        rowids: Optional[Iterator[int]] = None,
+        abort_on_exc: bool = False,
     ) -> tuple[int, int, set[int]]:
         """Insert rows into the store table and update the catalog table's md
         Returns:
@@ -347,12 +355,12 @@ class StoreBase:
 
                         if show_progress:
                             if progress_bar is None:
-                                warnings.simplefilter("ignore", category=TqdmWarning)
+                                warnings.simplefilter('ignore', category=TqdmWarning)
                                 progress_bar = tqdm(
                                     desc=f'Inserting rows into `{self.tbl_version.name}`',
                                     unit=' rows',
                                     ncols=100,
-                                    file=sys.stdout
+                                    file=sys.stdout,
                                 )
                             progress_bar.update(1)
 
@@ -379,8 +387,13 @@ class StoreBase:
         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.ColumnElement[bool]], conn: sql.engine.Connection) -> int:
+        self,
+        current_version: int,
+        base_versions: list[Optional[int]],
+        match_on_vmin: bool,
+        where_clause: Optional[sql.ColumnElement[bool]],
+        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:
@@ -394,12 +407,12 @@ class StoreBase:
         """
         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)
+            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)
+        base_versions_clause = (
+            sql.true() if len(base_versions) == 0 else self.base._versions_clause(base_versions, match_on_vmin)
+        )
         set_clause: dict[sql.Column, Union[int, sql.Column]] = {self.v_max_col: current_version}
         for index_info in self.tbl_version.idxs_by_name.values():
             # copy value column to undo column
@@ -450,7 +463,9 @@ class StoreView(StoreBase):
     def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.and_(
             self.base._rowid_join_predicate(),
-            *[c1 == c2 for c1, c2 in zip(self.rowid_columns(), self.base.rowid_columns())])
+            *[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
@@ -482,4 +497,5 @@ class StoreComponentView(StoreView):
     def _rowid_join_predicate(self) -> sql.ColumnElement[bool]:
         return sql.and_(
             self.base._rowid_join_predicate(),
-            *[c1 == c2 for c1, c2 in zip(self.rowid_columns()[:-1], self.base.rowid_columns())])
+            *[c1 == c2 for c1, c2 in zip(self.rowid_columns()[:-1], self.base.rowid_columns())],
+        )