pixeltable 0.3.14__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/share/packager.py

@@ -1,23 +1,35 @@
-import datetime
+import base64
+import io
 import json
 import logging
 import tarfile
 import urllib.parse
 import urllib.request
 import uuid
+from datetime import timedelta
 from pathlib import Path
-from typing import Any, Iterator, Optional
+from typing import Any, Iterator
+from uuid import UUID
 
 import more_itertools
+import numpy as np
+import pgvector.sqlalchemy as sql_vector  # type: ignore[import-untyped]
+import PIL.Image
 import pyarrow as pa
 import pyarrow.parquet as pq
 import sqlalchemy as sql
 
 import pixeltable as pxt
-from pixeltable import catalog, exceptions as excs, metadata
+import pixeltable.utils.av as av_utils
+from pixeltable import catalog, exceptions as excs, metadata, type_system as ts
+from pixeltable.catalog.table_version import TableVersionKey, TableVersionMd
 from pixeltable.env import Env
+from pixeltable.exprs.data_row import CellMd
 from pixeltable.metadata import schema
-from pixeltable.utils.media_store import MediaStore
+from pixeltable.utils import sha256sum
+from pixeltable.utils.formatter import Formatter
+from pixeltable.utils.local_store import TempStore
+from pixeltable.utils.object_stores import ObjectOps
 
 _logger = logging.getLogger('pixeltable')
 
@@ -42,43 +54,56 @@ class TablePackager:
     tmp_dir: Path  # Temporary directory where the package will reside
     tables_dir: Path  # Directory where the Parquet tables will be written
     media_files: dict[Path, str]  # Mapping from local media file paths to their tarball names
-    md: dict[str, Any]
+    bundle_md: dict[str, Any]
 
-    def __init__(self, table: catalog.Table, additional_md: Optional[dict[str, Any]] = None) -> None:
+    bundle_path: Path
+    preview_header: dict[str, str]
+    preview: list[list[Any]]
+
+    def __init__(self, table: catalog.Table, additional_md: dict[str, Any] | None = None) -> None:
         self.table = table
-        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
-        # Load metadata
-        with Env.get().begin_xact():
+        # Load metadata and convert to JSON immediately
+        with catalog.Catalog.get().begin_xact(for_write=False):
             tbl_md = catalog.Catalog.get().load_replica_md(table)
-            self.md = {
+            self.bundle_md = {
                 'pxt_version': pxt.__version__,
                 'pxt_md_version': metadata.VERSION,
-                'md': {'tables': [md.as_dict() for md in tbl_md]},
+                'md': [md.as_dict() for md in tbl_md],
             }
         if additional_md is not None:
-            self.md.update(additional_md)
+            self.bundle_md.update(additional_md)
 
     def package(self) -> Path:
         """
         Export the table to a tarball containing Parquet 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}")
+
+        _logger.info(f'Packaging table {self.table._path()!r} and its ancestors in: {self.tmp_dir}')
         self.tmp_dir.mkdir()
         with open(self.tmp_dir / 'metadata.json', 'w', encoding='utf8') as fp:
-            json.dump(self.md, fp)
+            json.dump(self.bundle_md, fp)
         self.tables_dir = self.tmp_dir / 'tables'
         self.tables_dir.mkdir()
-        with Env.get().begin_xact():
+        with catalog.Catalog.get().begin_xact(for_write=False):
             for tv in self.table._tbl_version_path.get_tbl_versions():
-                _logger.info(f"Exporting table '{tv.get().name}:{tv.get().version}'.")
+                _logger.info(f'Exporting table {tv.get().versioned_name!r}.')
                 self.__export_table(tv.get())
+
         _logger.info('Building archive.')
-        bundle_path = self.__build_tarball()
-        _logger.info(f'Packaging complete: {bundle_path}')
-        return bundle_path
+        self.bundle_path = self.__build_tarball()
+
+        _logger.info('Extracting preview data.')
+        self.bundle_md['row_count'] = self.table.count()
+        preview_header, preview = self.__extract_preview_data()
+        self.bundle_md['preview_header'] = preview_header
+        self.bundle_md['preview_data'] = preview
+
+        _logger.info(f'Packaging complete: {self.bundle_path}')
+        return self.bundle_path
 
     def __export_table(self, tv: catalog.TableVersion) -> None:
         """
@@ -88,9 +113,12 @@ class TablePackager:
         assert any(tv.id == base.id for base in self.table._tbl_version_path.get_tbl_versions())
         sql_types = {col.name: col.type for col in tv.store_tbl.sa_tbl.columns}
         media_cols: set[str] = set()
-        for col in tv.cols_by_name.values():
+        cellmd_cols: set[str] = set()
+        for col in tv.cols:
             if col.is_stored and col.col_type.is_media_type():
                 media_cols.add(col.store_name())
+            if col.stores_cellmd:
+                cellmd_cols.add(col.cellmd_store_name())
 
         parquet_schema = self.__to_parquet_schema(tv.store_tbl.sa_tbl)
         # TODO: Partition larger tables into multiple parquet files. (The parquet file naming scheme anticipates
@@ -101,14 +129,14 @@ class TablePackager:
         _logger.info(f'Creating parquet table: {parquet_file}')
 
         # Populate the Parquet table with data.
-        # The data is first loaded from the DataFrame into a sequence of pyarrow tables, batched in order to avoid
+        # The data is first loaded from the Query into a sequence of pyarrow tables, batched in order to avoid
         # excessive memory usage. The pyarrow tables are then amalgamated into the (single) Parquet table on disk.
         # We use snappy compression for the Parquet tables; the entire bundle will be bzip2-compressed later, so
         # faster compression should provide good performance while still reducing temporary storage utilization.
-        parquet_writer = pq.ParquetWriter(parquet_file, parquet_schema, compression='SNAPPY')
-        filter_tv = self.table._tbl_version.get()
+        parquet_writer = pq.ParquetWriter(parquet_file, parquet_schema, compression='snappy')
+        filter_tv = self.table._tbl_version_path.tbl_version.get()
         row_iter = tv.store_tbl.dump_rows(tv.version, filter_tv.store_tbl, filter_tv.version)
-        for pa_table in self.__to_pa_tables(row_iter, sql_types, media_cols, parquet_schema):
+        for pa_table in self.__to_pa_tables(row_iter, sql_types, media_cols, cellmd_cols, parquet_schema):
             parquet_writer.write_table(pa_table)
         parquet_writer.close()
 
@@ -117,7 +145,7 @@ class TablePackager:
     @classmethod
     def __to_parquet_schema(cls, store_tbl: sql.Table) -> pa.Schema:
         entries = [(col_name, cls.__to_parquet_type(col.type)) for col_name, col in store_tbl.columns.items()]
-        return pa.schema(entries)  # type: ignore[arg-type]
+        return pa.schema(entries)
 
     @classmethod
     def __to_parquet_type(cls, col_type: sql.types.TypeEngine[Any]) -> pa.DataType:
@@ -130,13 +158,19 @@ class TablePackager:
         if isinstance(col_type, sql.Float):
             return pa.float32()
         if isinstance(col_type, sql.TIMESTAMP):
-            return pa.timestamp('us', tz=datetime.timezone.utc)
+            return pa.timestamp('us', tz='UTC')
         if isinstance(col_type, sql.Date):
             return pa.date32()
         if isinstance(col_type, sql.JSON):
             return pa.string()  # JSON will be exported as strings
         if isinstance(col_type, sql.LargeBinary):
             return pa.binary()
+        if isinstance(col_type, sql.UUID):
+            return pa.uuid()
+        if isinstance(col_type, sql_vector.Vector):
+            # Parquet/pyarrow do not handle null values properly for fixed_shape_tensor(), so we have to use list_()
+            # here instead.
+            return pa.list_(pa.float32())
         raise AssertionError(f'Unrecognized SQL type: {col_type} (type {type(col_type)})')
 
     def __to_pa_tables(
@@ -144,6 +178,7 @@ class TablePackager:
         row_iter: Iterator[dict[str, Any]],
         sql_types: dict[str, sql.types.TypeEngine[Any]],
         media_cols: set[str],
+        cellmd_cols: set[str],
         arrow_schema: pa.Schema,
         batch_size: int = 1_000,
     ) -> Iterator[pa.Table]:
@@ -155,17 +190,28 @@ class TablePackager:
         for rows in more_itertools.batched(row_iter, batch_size):
             cols = {}
             for name, sql_type in sql_types.items():
-                is_media_col = name in media_cols
-                values = [self.__to_pa_value(row.get(name), sql_type, is_media_col) for row in rows]
+                values = [
+                    self.__to_pa_value(row.get(name), sql_type, name in media_cols, name in cellmd_cols) for row in rows
+                ]
                 cols[name] = values
             yield pa.Table.from_pydict(cols, schema=arrow_schema)
 
-    def __to_pa_value(self, val: Any, sql_type: sql.types.TypeEngine[Any], is_media_col: bool) -> Any:
+    def __to_pa_value(
+        self, val: Any, sql_type: sql.types.TypeEngine[Any], is_media_col: bool, is_cellmd_col: bool
+    ) -> Any:
         if val is None:
             return None
+        if is_cellmd_col:
+            assert isinstance(val, dict)
+            # Export JSON as strings
+            return json.dumps(self.__process_cellmd(val))
         if isinstance(sql_type, sql.JSON):
             # Export JSON as strings
             return json.dumps(val)
+        if isinstance(sql_type, sql.UUID):
+            # PyArrow's pa.uuid() expects bytes
+            assert isinstance(val, uuid.UUID)
+            return val.bytes
         if is_media_col:
             # Handle media files as described above
             assert isinstance(val, str)
@@ -173,6 +219,10 @@ class TablePackager:
         return val
 
     def __process_media_url(self, url: str) -> str:
+        """
+        Process a media URL for export. If it's a local file URL (file://), then replace it with a pxtmedia:// URI,
+        copying the file into the tarball if necessary. If it's any other type of URL, return it unchanged.
+        """
         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.
@@ -182,12 +232,32 @@ class TablePackager:
             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}'
+                # We name the media files in the archive by their SHA256 hash. This ensures that we can properly
+                # deduplicate and validate them later.
+                # If we get a collision, it's not a problem; it just means we have two identical files (which will
+                # be conveniently deduplicated in the bundle).
+                sha = sha256sum(path)
+                dest_name = f'{sha}{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 __process_cellmd(self, cellmd: dict[str, Any]) -> dict[str, Any]:
+        """
+        Process a cellmd dictionary for export. This involves replacing any local file references
+        with pxtmedia:// URIs, as described above.
+        """
+        cellmd_ = CellMd.from_dict(cellmd)
+        if cellmd_.file_urls is None:
+            return cellmd  # No changes
+
+        updated_urls: list[str] = []
+        for url in cellmd_.file_urls:
+            updated_urls.append(self.__process_media_url(url))
+        cellmd_.file_urls = updated_urls
+        return cellmd_.as_dict()
+
     def __build_tarball(self) -> Path:
         bundle_path = self.tmp_dir / 'bundle.tar.bz2'
         with tarfile.open(bundle_path, 'w:bz2') as tf:
@@ -200,6 +270,116 @@ class TablePackager:
                 tf.add(src_file, arcname=f'media/{dest_name}')
         return bundle_path
 
+    def __extract_preview_data(self) -> tuple[dict[str, str], list[list[Any]]]:
+        """
+        Extract a preview of the table data for display in the UI.
+
+        In order to bound the size of the output data, all "unbounded" data types are resized:
+        - Strings are abbreviated as per Formatter.abbreviate()
+        - Arrays and JSON are shortened and formatted as strings
+        - Images are resized to thumbnail size as a base64-encoded webp
+        - Videos are replaced by their first frame and resized as above
+        - Documents are replaced by a thumbnail as a base64-encoded webp
+        """
+        preview_cols = self.table._get_schema()
+        select_list = [self.table[col_name] for col_name in preview_cols]
+        # First 5 rows
+        rows = list(self.table.select(*select_list).head(n=5))
+
+        preview_header = {col_name: str(col_type._type) for col_name, col_type in preview_cols.items()}
+        preview = [
+            [self.__encode_preview_data(val, col_type)]
+            for row in rows
+            for val, col_type in zip(row.values(), preview_cols.values(), strict=True)
+        ]
+
+        return preview_header, preview
+
+    def __encode_preview_data(self, val: Any, col_type: ts.ColumnType) -> Any:
+        if val is None:
+            return None
+
+        match col_type._type:
+            case ts.ColumnType.Type.STRING:
+                assert isinstance(val, str)
+                return Formatter.abbreviate(val)
+
+            case ts.ColumnType.Type.INT | ts.ColumnType.Type.FLOAT | ts.ColumnType.Type.BOOL:
+                return val
+
+            case ts.ColumnType.Type.TIMESTAMP | ts.ColumnType.Type.DATE:
+                return str(val)
+
+            case ts.ColumnType.Type.UUID:
+                assert isinstance(val, uuid.UUID)
+                return str(val)
+
+            case ts.ColumnType.Type.BINARY:
+                assert isinstance(val, bytes)
+                return Formatter.format_binary(val)
+
+            case ts.ColumnType.Type.ARRAY:
+                assert isinstance(val, np.ndarray)
+                return Formatter.format_array(val)
+
+            case ts.ColumnType.Type.JSON:
+                # We need to escape the JSON string server-side for security reasons.
+                # Therefore we don't escape it here, in order to avoid double-escaping.
+                return Formatter.format_json(val, escape_strings=False)
+
+            case ts.ColumnType.Type.IMAGE:
+                # Rescale the image to minimize data transfer size
+                assert isinstance(val, PIL.Image.Image)
+                return self.__encode_image(val)
+
+            case ts.ColumnType.Type.AUDIO:
+                assert isinstance(val, str)
+                return self.__encode_audio(val)
+
+            case ts.ColumnType.Type.VIDEO:
+                assert isinstance(val, str)
+                return self.__encode_video(val)
+
+            case ts.ColumnType.Type.DOCUMENT:
+                assert isinstance(val, str)
+                return self.__encode_document(val)
+
+            case _:
+                raise AssertionError(f'Unrecognized column type: {col_type._type}')
+
+    def __encode_image(self, img: PIL.Image.Image) -> str:
+        # Heuristic for thumbnail sizing:
+        # Standardize on a width of 240 pixels (to most efficiently utilize the columnar display).
+        # But, if the aspect ratio is below 2:3, bound the height at 360 pixels (to avoid unboundedly tall thumbnails
+        #     in the case of highly oblong images).
+        if img.height > img.width * 1.5:
+            scaled_img = img.resize((img.width * 360 // img.height, 360))
+        else:
+            scaled_img = img.resize((240, img.height * 240 // img.width))
+        with io.BytesIO() as buffer:
+            scaled_img.save(buffer, 'webp')
+            return base64.b64encode(buffer.getvalue()).decode()
+
+    def __encode_audio(self, audio_path: str) -> str | None:
+        try:
+            audio_md = av_utils.get_metadata(audio_path)
+            if 'streams' in audio_md:
+                duration = audio_md['streams'][0]['duration_seconds']
+                assert isinstance(duration, float)
+                return f'{timedelta(seconds=round(duration))} audio clip'
+            return None
+        except Exception:
+            _logger.info(f'Could not extract audio metadata from file for data preview: {audio_path}', exc_info=True)
+            return None
+
+    def __encode_video(self, video_path: str) -> str | None:
+        thumb = Formatter.extract_first_video_frame(video_path)
+        return self.__encode_image(thumb) if thumb is not None else None
+
+    def __encode_document(self, doc_path: str) -> str | None:
+        thumb = Formatter.make_document_thumbnail(doc_path)
+        return self.__encode_image(thumb) if thumb is not None else None
+
 
 class TableRestorer:
     """
@@ -208,34 +388,35 @@ class TableRestorer:
 
     Args:
         tbl_path: Pixeltable path (such as 'my_dir.my_table') where the materialized table will be made visible.
-        md: Optional metadata dictionary. If not provided, metadata will be read from the tarball's `metadata.json`.
+        bundle_md: Optional metadata dictionary.
+            If not provided, metadata will be read from the tarball's `metadata.json`.
             The metadata contains table_md, table_version_md, and table_schema_version_md entries for each ancestor
             of the table being restored, as written out by `TablePackager`.
     """
 
     tbl_path: str
-    md: Optional[dict[str, Any]]
+    bundle_md: dict[str, Any] | None
     tmp_dir: Path
     media_files: dict[str, str]  # Mapping from pxtmedia:// URLs to local file:// URLs
 
-    def __init__(self, tbl_path: str, md: Optional[dict[str, Any]] = None) -> None:
+    def __init__(self, tbl_path: str, bundle_md: dict[str, Any] | None = None) -> None:
         self.tbl_path = tbl_path
-        self.md = md
-        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.bundle_md = bundle_md
+        self.tmp_dir = TempStore.create_path()
         self.media_files = {}
 
-    def restore(self, bundle_path: Path) -> pxt.Table:
+    def restore(self, bundle_path: Path, pxt_uri: str | None = None, explicit_version: int | None = None) -> pxt.Table:
         # Extract tarball
         print(f'Extracting table data into: {self.tmp_dir}')
         with tarfile.open(bundle_path, 'r:bz2') as tf:
             tf.extractall(path=self.tmp_dir)
 
-        if self.md is None:
+        if self.bundle_md is None:
             # No metadata supplied; read it from the archive
             with open(self.tmp_dir / 'metadata.json', 'r', encoding='utf8') as fp:
-                self.md = json.load(fp)
+                self.bundle_md = json.load(fp)
 
-        pxt_md_version = self.md['pxt_md_version']
+        pxt_md_version = self.bundle_md['pxt_md_version']
         assert isinstance(pxt_md_version, int)
 
         if pxt_md_version != metadata.VERSION:
@@ -243,44 +424,301 @@ class TableRestorer:
                 f'Pixeltable metadata version mismatch: {pxt_md_version} != {metadata.VERSION}.\n'
                 'Please upgrade Pixeltable to use this dataset: pip install -U pixeltable'
             )
+        # Convert tables metadata from dict to list of TableVersionMd
+        tbl_md = [schema.md_from_dict(TableVersionMd, t) for t in self.bundle_md['md']]
 
-        tbl_md = [schema.FullTableMd.from_dict(t) for t in self.md['md']['tables']]
+        for md in tbl_md:
+            md.tbl_md.is_replica = True
 
-        # Create the replica table
-        # TODO: This needs to be made concurrency-safe.
-        replica_tbl = catalog.Catalog.get().create_replica(catalog.Path(self.tbl_path), tbl_md)
-        assert replica_tbl._tbl_version.get().is_snapshot
+        assert not tbl_md[0].version_md.is_fragment  # Top-level table cannot be a version fragment
 
-        # Now we need to instantiate and load data for replica_tbl and its ancestors, except that we skip
-        # replica_tbl itself if it's a pure snapshot.
-        if replica_tbl._id != replica_tbl._tbl_version.id:
-            ancestor_md = tbl_md[1:]  # Pure snapshot; skip replica_tbl
-        else:
-            ancestor_md = tbl_md  # Not a pure snapshot; include replica_tbl
+        cat = catalog.Catalog.get()
 
-        # Instantiate data from the Parquet tables.
-        with Env.get().begin_xact():
-            for md in ancestor_md[::-1]:  # Base table first
-                # Create a TableVersion instance (and a store table) for this ancestor.
-                tv = catalog.TableVersion.create_replica(md)
-                # Now import data from Parquet.
-                _logger.info(f'Importing table {tv.name!r}.')
-                self.__import_table(self.tmp_dir, tv, md)
+        with cat.begin_xact(for_write=True):
+            # Create (or update) the replica table and its ancestors, along with TableVersion instances for any
+            # versions that have not been seen before.
+            cat.create_replica(catalog.Path.parse(self.tbl_path), tbl_md)
 
-        return replica_tbl
+            _logger.debug(f'Now will import data for {len(tbl_md)} table(s):')
+            _logger.debug(repr([md.tbl_md.tbl_id for md in tbl_md[::-1]]))
 
-    def __import_table(self, bundle_path: Path, tv: catalog.TableVersion, tbl_md: schema.FullTableMd) -> None:
+            # Now we need to load data for replica_tbl and its ancestors, except that we skip
+            # replica_tbl itself if it's a pure snapshot.
+            for md in tbl_md[::-1]:  # Base table first
+                if not md.is_pure_snapshot:
+                    tv = cat.get_tbl_version(TableVersionKey(UUID(md.tbl_md.tbl_id), md.version_md.version, None))
+                    # Import data from Parquet.
+                    _logger.info(f'Importing table {tv.name!r}.')
+                    self.__import_table(self.tmp_dir, tv, md)
+
+            tbl = cat.get_table_by_id(UUID(tbl_md[0].tbl_md.tbl_id), version=explicit_version)
+            if pxt_uri is not None:
+                # Set pxt_uri for the newly created table
+                cat.update_additional_md(tbl._id, {'pxt_uri': pxt_uri})
+            tbl._tbl_version_path.clear_cached_md()  # TODO: Clear cached md for ancestors too?
+            return tbl
+
+    def __import_table(self, bundle_path: Path, tv: catalog.TableVersion, tbl_md: TableVersionMd) -> None:
         """
         Import the Parquet table into the Pixeltable catalog.
         """
-        tbl_id = uuid.UUID(tbl_md.tbl_md.tbl_id)
+        tbl_id = UUID(tbl_md.tbl_md.tbl_id)
         parquet_dir = bundle_path / 'tables' / f'tbl_{tbl_id.hex}'
         parquet_table = pq.read_table(str(parquet_dir))
-
-        for batch in parquet_table.to_batches():
+        replica_version = tv.version
+
+        conn = Env.get().conn
+        store_sa_tbl = tv.store_tbl.sa_tbl
+        store_sa_tbl_name = tv.store_tbl._storage_name()
+
+        # Sometimes we are importing a table that has never been seen before. Other times, however, we are importing
+        # an existing replica table, and the table version and/or row selection differs from what was imported
+        # previously. Care must be taken to ensure that the new data is merged with existing data in a way that
+        # yields an internally consistent version history for each row.
+
+        # The overall strategy is this:
+        # 1. Import the parquet data into a temporary table;
+        # 2. "rectify" the v_max values in both the temporary table and the existing table (more on this below);
+        # 3. Delete any row instances from the temporary table that are already present in the existing table;
+        # 4. Copy the remaining rows from the temporary table into the existing table.
+        # 5. Rectify any index columns.
+
+        # STEP 1: Import the parquet data into a temporary table.
+
+        # Create a temporary table for the initial data load, containing columns for all columns present in the
+        # parquet table. The parquet columns have identical names to those in the store table, so we can use the
+        # store table schema to get their SQL types (which are not necessarily derivable from their Parquet types,
+        # e.g., pa.string() may hold either VARCHAR or serialized JSONB).
+        temp_cols: dict[str, sql.Column] = {}
+        for field in parquet_table.schema:
+            assert field.name in store_sa_tbl.columns, f'{field.name} not in {list(store_sa_tbl.columns)}'
+            col_type = store_sa_tbl.columns[field.name].type
+            temp_cols[field.name] = sql.Column(field.name, col_type)
+        temp_sa_tbl_name = f'temp_{uuid.uuid4().hex}'
+        _logger.debug(f'Creating temporary table: {temp_sa_tbl_name}')
+        temp_md = sql.MetaData()
+        temp_sa_tbl = sql.Table(temp_sa_tbl_name, temp_md, *temp_cols.values(), prefixes=['TEMPORARY'])
+        temp_sa_tbl.create(conn)
+
+        # Populate the temporary table with data from the Parquet file.
+        _logger.debug(f'Loading {parquet_table.num_rows} row(s) into temporary table: {temp_sa_tbl_name}')
+        for batch in parquet_table.to_batches(max_chunksize=10_000):
             pydict = batch.to_pydict()
             rows = self.__from_pa_pydict(tv, pydict)
-            tv.store_tbl.load_rows(rows)
+            conn.execute(sql.insert(temp_sa_tbl), rows)
+
+        # STEP 2: Rectify v_max values.
+
+        # Each row version is identified uniquely by its pk, a tuple (row_id, pos_0, pos_1, ..., pos_k, v_min).
+        # Conversely, v_max is not part of the primary key, but is simply a bookkeeping device.
+        # In an original table, v_max is always equal to the v_min of the succeeding row instance with the same
+        # row id, or MAX_VERSION if no such row instance exists. But in the replica, we need to be careful, since
+        # we might see only a subset of the original table's versions, and we might see them out of order.
+
+        # We'll adjust the v_max values according to the principle of "latest provable v_max":
+        # they will always correspond to the latest version for which we can prove the row instance was alive. This
+        # will enable us to maintain consistency of the v_max values if additional table versions are later imported,
+        # regardless of the order in which they are seen. It also means that replica tables (unlike original tables)
+        # may have gaps in their row version histories, but this is fine; the gaps simply correspond to table versions
+        # that have never been observed.
+
+        pk_predicates = [col == temp_cols[col.name] for col in tv.store_tbl.pk_columns()]
+        pk_clause = sql.and_(*pk_predicates)
+
+        # If the same pk exists in both the temporary table and the existing table, then the corresponding row data
+        # must be identical; the rows can differ only in their v_max value. As a sanity check, we go through the
+        # motion of verifying this; a failure implies data corruption in either the replica being imported or in a
+        # previously imported replica.
+
+        system_col_names = {col.name for col in tv.store_tbl.system_columns()}
+        media_col_names = {col.store_name() for col in tv.cols if col.col_type.is_media_type() and col.is_stored}
+        value_store_cols = [
+            store_sa_tbl.c[col_name]
+            for col_name in temp_cols
+            if col_name not in system_col_names and col_name not in media_col_names
+        ]
+        value_temp_cols = [
+            col
+            for col_name, col in temp_cols.items()
+            if col_name not in system_col_names and col_name not in media_col_names
+        ]
+
+        q: sql.Executable
+
+        assert len(value_store_cols) == len(value_temp_cols)
+        if len(value_store_cols) > 0:
+            mismatch_predicates = [
+                store_col != temp_col for store_col, temp_col in zip(value_store_cols, value_temp_cols)
+            ]
+            mismatch_clause = sql.or_(*mismatch_predicates)
+
+            # This query looks for rows that have matching primary keys (rowid + pos_k + v_min), but differ in at least
+            # one value column. Pseudo-SQL:
+            #
+            # SELECT store_tbl.col_0, ..., store_tbl.col_n, temp_tbl.col_0, ...,  temp_tbl.col_n
+            # FROM store_tbl, temp_tbl
+            # WHERE store_tbl.rowid = temp_tbl.rowid
+            #     AND store_tbl.pos_0 = temp_tbl.pos_0
+            #     AND ... AND store_tbl.pos_k = temp_tbl.pos_k
+            #     AND store_tbl.v_min = temp_tbl.v_min
+            #     AND (
+            #         store_tbl.col_0 != temp_tbl.col_0
+            #         OR store_tbl.col_1 != temp_tbl.col_1
+            #         OR ... OR store_tbl.col_n != temp_tbl.col_n
+            #     )
+            #
+            # The value column comparisons (store_tbl.col_0 != temp_tbl.col_0, etc.) will always be false for rows where
+            # either column is NULL; this is what we want, since it may indicate a column that is present in one version
+            # but not the other.
+            q = sql.select(*value_store_cols, *value_temp_cols).where(pk_clause).where(mismatch_clause)
+            _logger.debug(q.compile())
+            result = conn.execute(q)
+            if result.rowcount > 0:
+                _logger.debug(
+                    f'Data corruption error between {temp_sa_tbl_name!r} and {store_sa_tbl_name!r}: '
+                    f'{result.rowcount} inconsistent row(s).'
+                )
+                row = result.first()
+                _logger.debug('Example mismatch:')
+                _logger.debug(f'{store_sa_tbl_name}: {row[: len(value_store_cols)]}')
+                _logger.debug(f'{temp_sa_tbl_name}: {row[len(value_store_cols) :]}')
+                raise excs.Error(
+                    'Data corruption error: '
+                    'the replica data are inconsistent with data retrieved from a previous replica.'
+                )
+
+        _logger.debug(f'Verified data integrity between {store_sa_tbl_name!r} and {temp_sa_tbl_name!r}.')
+
+        # Now rectify the v_max values in the temporary table.
+        # If a row instance has a concrete v_max value, then we know it's genuine: it's the unique and immutable
+        # version when the row was deleted. (This can only happen if later versions of the base table already
+        # existed at the time this replica was published.)
+        # But if a row instance has a v_max value of MAX_VERSION, then we don't know anything about its future.
+        # It might live indefinitely, or it might be deleted as early as version `n + 1`. Following the principle
+        # of "latest provable v_max", we simply set v_max equal to `n + 1`.
+        q = (
+            temp_sa_tbl.update()
+            .values(v_max=(replica_version + 1))
+            .where(temp_sa_tbl.c.v_max == schema.Table.MAX_VERSION)
+        )
+        _logger.debug(q.compile())
+        result = conn.execute(q)
+        _logger.debug(f'Rectified {result.rowcount} row(s) in {temp_sa_tbl_name!r}.')
+
+        # Now rectify the v_max values in the existing table. This is done by simply taking the later of the two v_max
+        # values (the existing one and the new one) for each row instance, following the "latest provable v_max"
+        # principle. Obviously we only need to do this for rows that exist in both tables (it's a simple join).
+        q = (
+            store_sa_tbl.update()
+            .values(v_max=sql.func.greatest(store_sa_tbl.c.v_max, temp_sa_tbl.c.v_max))
+            .where(pk_clause)
+        )
+        _logger.debug(q.compile())
+        result = conn.execute(q)
+        _logger.debug(f'Rectified {result.rowcount} row(s) in {store_sa_tbl_name!r}.')
+
+        # STEP 3: Delete any row instances from the temporary table that are already present in the existing table.
+
+        # Now we need to update rows in the existing table that are also present in the temporary table. This is to
+        # account for the scenario where the temporary table has columns that are not present in the existing table.
+        # (We can't simply replace the rows with their versions in the temporary table, because the converse scenario
+        # might also occur; there may be columns in the existing table that are not present in the temporary table.)
+        value_update_clauses: dict[str, sql.ColumnElement] = {}
+        for temp_col in temp_cols.values():
+            if temp_col.name not in system_col_names:
+                store_col = store_sa_tbl.c[temp_col.name]
+                # Prefer the value from the existing table, substituting the value from the temporary table if it's
+                # NULL. This works in all cases (including media columns, where we prefer the existing media file).
+                clause = sql.case((store_col == None, temp_col), else_=store_col)
+                value_update_clauses[temp_col.name] = clause
+        if len(value_update_clauses) > 0:
+            q = store_sa_tbl.update().values(**value_update_clauses).where(pk_clause)
+            _logger.debug(q.compile())
+            result = conn.execute(q)
+            _logger.debug(
+                f'Merged values from {temp_sa_tbl_name!r} into {store_sa_tbl_name!r} for {result.rowcount} row(s).'
+            )
+
+        # Now drop any rows from the temporary table that are also present in the existing table.
+        # The v_max values have been rectified, data has been merged into NULL cells, and all other row values have
+        # been verified identical.
+        # TODO: Delete any media files that were orphaned by this operation (they're necessarily duplicates of media
+        #     files that are already present in the existing table).
+        q = temp_sa_tbl.delete().where(pk_clause)
+        _logger.debug(q.compile())
+        result = conn.execute(q)
+        _logger.debug(f'Deleted {result.rowcount} row(s) from {temp_sa_tbl_name!r}.')
+
+        # STEP 4: Copy the remaining rows from the temporary table into the existing table.
+
+        # Now copy the remaining data (consisting entirely of new row instances) from the temporary table into
+        # the actual table.
+        q = store_sa_tbl.insert().from_select(
+            [store_sa_tbl.c[col_name] for col_name in temp_cols], sql.select(*temp_cols.values())
+        )
+        _logger.debug(q.compile())
+        result = conn.execute(q)
+        _logger.debug(f'Inserted {result.rowcount} row(s) from {temp_sa_tbl_name!r} into {store_sa_tbl_name!r}.')
+
+        # STEP 5: Rectify any index columns.
+
+        # Finally, rectify any index columns in the table. This involves shuffling data between the index's val and
+        # undo columns to ensure they appropriately reflect the most recent replicated version of the table.
+
+        # Get the most recent replicated version of the table. This might be the version we're currently importing,
+        # but it might be a different version of the table that was previously imported.
+        head_version_md = catalog.Catalog.get()._collect_tbl_history(tv.id, n=1)[0]
+        head_version = head_version_md.version_md.version
+        _logger.debug(f'Head version for index rectification is {head_version}.')
+
+        # Get the index info from the table metadata. Here we use the tbl_md that we just collected from the DB.
+        # This is to ensure we pick up ALL indices, including dropped indices and indices that are present in
+        # a previously replicated version of the table, but not in the one currently being imported.
+        index_md = head_version_md.tbl_md.index_md
+
+        # Now update the table. We can do this for all indices together with just two SQL queries. For each index,
+        # at most one of the val or undo columns will be non-NULL in any given row.
+        # For rows where v_min <= head_version < v_max, we set, for all indices:
+        #   val_col = whichever of (val_col, undo_col) is non-NULL (or NULL if both are, e.g., for a dropped index)
+        #   undo_col = NULL
+        # For rows where head_version < v_min or v_max <= head_version, vice versa.
+        val_sql_clauses: dict[str, sql.ColumnElement] = {}
+        undo_sql_clauses: dict[str, sql.ColumnElement] = {}
+        for index in index_md.values():
+            if index.class_fqn.endswith('.EmbeddingIndex'):
+                val_col_name = f'col_{index.index_val_col_id}'
+                undo_col_name = f'col_{index.index_val_undo_col_id}'
+                # Check that the val column for the index is actually present in the store table. We need to do this
+                # to properly handle the case where the replica represents a table version that was *not* the most
+                # recent version at the time it was published. In that case, it is possible for tbl_md to contain
+                # metadata for indices not known to any version that has been replicated. (However, the converse
+                # *does* hold: all replicated indices must have metadata in tbl_md; and that's what's important.)
+                if val_col_name in store_sa_tbl.c:
+                    assert undo_col_name in store_sa_tbl.c
+                    coalesce = sql.func.coalesce(store_sa_tbl.c[val_col_name], store_sa_tbl.c[undo_col_name])
+                    val_sql_clauses[val_col_name] = coalesce
+                    val_sql_clauses[undo_col_name] = sql.null()
+                    undo_sql_clauses[undo_col_name] = coalesce
+                    undo_sql_clauses[val_col_name] = sql.null()
+
+        if len(val_sql_clauses) > 0:
+            q2 = (
+                store_sa_tbl.update()
+                .values(**val_sql_clauses)
+                .where(sql.and_(tv.store_tbl.v_min_col <= head_version, tv.store_tbl.v_max_col > head_version))
+            )
+            _logger.debug(q2.compile())
+            _ = conn.execute(q2)
+            q2 = (
+                store_sa_tbl.update()
+                .values(**undo_sql_clauses)
+                .where(sql.or_(tv.store_tbl.v_min_col > head_version, tv.store_tbl.v_max_col <= head_version))
+            )
+            _logger.debug(q2.compile())
+            _ = conn.execute(q2)
+            _logger.debug(f'Rectified index columns in {store_sa_tbl_name!r}.')
+        else:
+            _logger.debug(f'No index columns to rectify in {store_sa_tbl_name!r}.')
 
     def __from_pa_pydict(self, tv: catalog.TableVersion, pydict: dict[str, Any]) -> list[dict[str, Any]]:
         # Data conversions from pyarrow to Pixeltable
@@ -288,36 +726,52 @@ class TableRestorer:
         for col_name in pydict:
             assert col_name in tv.store_tbl.sa_tbl.columns
             sql_types[col_name] = tv.store_tbl.sa_tbl.columns[col_name].type
-        media_col_ids: dict[str, int] = {}
-        for col in tv.cols_by_name.values():
-            if col.is_stored and col.col_type.is_media_type():
-                media_col_ids[col.store_name()] = col.id
+        stored_cols: dict[str, catalog.Column] = {col.store_name(): col for col in tv.cols if col.is_stored}
+        stored_cols |= {col.cellmd_store_name(): col for col in tv.cols if col.stores_cellmd}
 
         row_count = len(next(iter(pydict.values())))
-        rows: list[dict[str, Any]] = []
-        for i in range(row_count):
-            row = {
-                col_name: self.__from_pa_value(tv, col_vals[i], sql_types[col_name], media_col_ids.get(col_name))
-                for col_name, col_vals in pydict.items()
-            }
-            rows.append(row)
+        rows: list[dict[str, Any]] = [{} for _ in range(row_count)]
+        for col_name, col_vals in pydict.items():
+            assert len(col_vals) == row_count
+            col = stored_cols.get(col_name)  # Will be None for system columns
+            is_media_col = col is not None and col.is_stored and col.col_type.is_media_type()
+            is_cellmd_col = col is not None and col.stores_cellmd and col_name == col.cellmd_store_name()
+            assert col is None or is_cellmd_col or col_name == col.store_name()
+
+            for i, val in enumerate(col_vals):
+                rows[i][col_name] = self.__from_pa_value(val, sql_types[col_name], col, is_media_col, is_cellmd_col)
 
         return rows
 
     def __from_pa_value(
-        self, tv: catalog.TableVersion, val: Any, sql_type: sql.types.TypeEngine[Any], media_col_id: Optional[int]
+        self,
+        val: Any,
+        sql_type: sql.types.TypeEngine[Any],
+        col: catalog.Column | None,
+        is_media_col: bool,
+        is_cellmd_col: bool,
     ) -> Any:
         if val is None:
             return None
+        if isinstance(sql_type, sql_vector.Vector):
+            if isinstance(val, list):
+                val = np.array(val, dtype=np.float32)
+            assert isinstance(val, np.ndarray) and val.dtype == np.float32 and val.ndim == 1
+            return val
+        if is_cellmd_col:
+            assert col is not None
+            assert isinstance(val, str)
+            return self.__restore_cellmd(col, json.loads(val))
         if isinstance(sql_type, sql.JSON):
             return json.loads(val)
-        if media_col_id is not None:
-            assert isinstance(val, str)
-            return self.__relocate_media_file(tv, media_col_id, val)
+        if is_media_col:
+            assert col is not None
+            return self.__relocate_media_file(col, val)
         return val
 
-    def __relocate_media_file(self, tv: catalog.TableVersion, media_col_id: int, url: str) -> str:
+    def __relocate_media_file(self, media_col: catalog.Column, url: str) -> str:
         # If this is a pxtmedia:// URL, relocate it
+        assert isinstance(url, str)
         parsed_url = urllib.parse.urlparse(url)
         assert parsed_url.scheme != 'file'  # These should all have been converted to pxtmedia:// URLs
         if parsed_url.scheme == 'pxtmedia':
@@ -325,9 +779,19 @@ class TableRestorer:
                 # First time seeing this pxtmedia:// URL. Relocate the file to the media store and record the mapping
                 # in self.media_files.
                 src_path = self.tmp_dir / 'media' / parsed_url.netloc
-                dest_path = MediaStore.prepare_media_path(tv.id, media_col_id, tv.version, ext=src_path.suffix)
-                src_path.rename(dest_path)
-                self.media_files[url] = urllib.parse.urljoin('file:', urllib.request.pathname2url(str(dest_path)))
+                # Move the file to the media store and update the URL.
+                self.media_files[url] = ObjectOps.put_file(media_col, src_path, relocate_or_delete=True)
             return self.media_files[url]
         # For any type of URL other than a local file, just return the URL as-is.
         return url
+
+    def __restore_cellmd(self, col: catalog.Column, cellmd: dict[str, Any]) -> dict[str, Any]:
+        cellmd_ = CellMd.from_dict(cellmd)
+        if cellmd_.file_urls is None:
+            return cellmd  # No changes
+
+        updated_urls: list[str] = []
+        for url in cellmd_.file_urls:
+            updated_urls.append(self.__relocate_media_file(col, url))
+        cellmd_.file_urls = updated_urls
+        return cellmd_.as_dict()