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

pixeltable/share/packager.py

@@ -1,4 +1,4 @@
-import io
+import datetime
 import json
 import logging
 import tarfile
@@ -9,45 +9,38 @@ from pathlib import Path
 from typing import Any, Iterator, Optional
 
 import more_itertools
-import numpy as np
 import pyarrow as pa
-import pyiceberg.catalog
+import pyarrow.parquet as pq
+import sqlalchemy as sql
 
 import pixeltable as pxt
-import pixeltable.type_system as ts
-from pixeltable import catalog, exprs, metadata
-from pixeltable.dataframe import DataFrame
+from pixeltable import catalog, exceptions as excs, metadata
 from pixeltable.env import Env
-from pixeltable.utils.arrow import PXT_TO_PA_TYPES
-from pixeltable.utils.iceberg import sqlite_catalog
+from pixeltable.metadata import schema
+from pixeltable.utils.media_store import MediaStore
 
 _logger = logging.getLogger('pixeltable')
 
 
 class TablePackager:
     """
-    Packages a pixeltable Table into a tarball containing Iceberg tables and media files. The structure of the tarball
+    Packages a pixeltable Table into a tarball containing Parquet tables and media files. The structure of the tarball
     is as follows:
 
-    metadata.json  # Pixeltable metadata for the packaged table
-    warehouse/catalog.db  # sqlite Iceberg catalog
-    warehouse/pxt.db/**  # Iceberg metadata and data files (parquet/avro/json)
+    metadata.json  # Pixeltable metadata for the packaged table and its ancestors
+    tables/**  # Parquet tables for the packaged table and its ancestors, each table in a directory 'tbl_{tbl_id.hex}'
     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.
+      stored in the Parquet 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}'.
+      'media/{uuid}{extension}', and the Parquet table will contain the ephemeral URI 'pxtmedia://{uuid}{extension}'.
     """
 
     table: catalog.Table  # The table to be packaged
     tmp_dir: Path  # Temporary directory where the package will reside
-    iceberg_catalog: pyiceberg.catalog.Catalog
+    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]
 
@@ -69,138 +62,113 @@ class TablePackager:
 
     def package(self) -> Path:
         """
-        Export the table to a tarball containing Iceberg tables and media files.
+        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}")
         self.tmp_dir.mkdir()
         with open(self.tmp_dir / 'metadata.json', 'w', encoding='utf8') as fp:
             json.dump(self.md, fp)
-        self.iceberg_catalog = sqlite_catalog(self.tmp_dir / 'warehouse')
+        self.tables_dir = self.tmp_dir / 'tables'
+        self.tables_dir.mkdir()
         with Env.get().begin_xact():
-            ancestors = (self.table, *self.table._bases)
-            for t in ancestors:
-                _logger.info(f"Exporting table '{t._path}'.")
-                self.__export_table(t)
+            for tv in self.table._tbl_version_path.get_tbl_versions():
+                _logger.info(f"Exporting table '{tv.get().name}:{tv.get().version}'.")
+                self.__export_table(tv.get())
         _logger.info('Building archive.')
         bundle_path = self.__build_tarball()
         _logger.info(f'Packaging complete: {bundle_path}')
         return bundle_path
 
-    def __export_table(self, t: catalog.Table) -> None:
+    def __export_table(self, tv: catalog.TableVersion) -> None:
         """
-        Exports the data from `t` into an Iceberg table.
+        Exports the data from `t` into a Parquet 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.get().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)
+        # `tv` must be an ancestor of the primary table
+        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():
+            if col.is_stored and col.col_type.is_media_type():
+                media_cols.add(col.store_name())
 
-    @classmethod
-    def __iceberg_namespace(cls, table: catalog.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}'
+        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
+        #     future support for this.)
+        parquet_dir = self.tables_dir / f'tbl_{tv.id.hex}'
+        parquet_dir.mkdir()
+        parquet_file = parquet_dir / f'tbl_{tv.id.hex}.00000.parquet'
+        _logger.info(f'Creating parquet table: {parquet_file}')
 
-    # 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
+        # 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
+        # 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()
+        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):
+            parquet_writer.write_table(pa_table)
+        parquet_writer.close()
+
+    # The following methods are responsible for schema and data conversion from Pixeltable to Parquet.
 
     @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()))
+    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]
 
     @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():
+    def __to_parquet_type(cls, col_type: sql.types.TypeEngine[Any]) -> pa.DataType:
+        if isinstance(col_type, sql.String):
             return pa.string()
-        return PXT_TO_PA_TYPES.get(col_type.__class__)
+        if isinstance(col_type, sql.Boolean):
+            return pa.bool_()
+        if isinstance(col_type, sql.BigInteger):
+            return pa.int64()
+        if isinstance(col_type, sql.Float):
+            return pa.float32()
+        if isinstance(col_type, sql.TIMESTAMP):
+            return pa.timestamp('us', tz=datetime.timezone.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()
+        raise AssertionError(f'Unrecognized SQL type: {col_type} (type {type(col_type)})')
 
     def __to_pa_tables(
-        self, df: DataFrame, actual_col_types: list[ts.ColumnType], arrow_schema: pa.Schema, batch_size: int = 1_000
+        self,
+        row_iter: Iterator[dict[str, Any]],
+        sql_types: dict[str, sql.types.TypeEngine[Any]],
+        media_cols: set[str],
+        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.
+        Group rows into a sequence of pyarrow tables, batched into smaller chunks to minimize memory utilization.
+        The row dictionaries have the format {store_col_name: value}, where the values reflect the unprocessed contents
+        of the store database (as returned by `StoreTable.dump_rows()`).
         """
-        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]
+        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]
+                cols[name] = values
             yield pa.Table.from_pydict(cols, schema=arrow_schema)
 
-    def __to_pa_rows(self, df: DataFrame, actual_col_types: list[ts.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:
+    def __to_pa_value(self, val: Any, sql_type: sql.types.TypeEngine[Any], is_media_col: bool) -> 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():
+        if isinstance(sql_type, sql.JSON):
             # Export JSON as strings
             return json.dumps(val)
-        if col_type.is_media_type():
+        if is_media_col:
             # Handle media files as described above
-            assert isinstance(val, str)  # Media columns are always referenced by `fileurl`
+            assert isinstance(val, str)
             return self.__process_media_url(val)
         return val
 
@@ -225,9 +193,141 @@ class TablePackager:
         with tarfile.open(bundle_path, 'w:bz2') as tf:
             # Add metadata json
             tf.add(self.tmp_dir / 'metadata.json', arcname='metadata.json')
-            # Add the Iceberg warehouse dir (including the catalog)
-            tf.add(self.tmp_dir / 'warehouse', arcname='warehouse', recursive=True)
+            # Add the dir containing Parquet tables
+            tf.add(self.tables_dir, arcname='tables')
             # 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
+
+
+class TableRestorer:
+    """
+    Creates a replica table from a tarball containing Parquet tables and media files. See the `TablePackager` docs for
+    details on the tarball structure.
+
+    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`.
+            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]]
+    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:
+        self.tbl_path = tbl_path
+        self.md = md
+        self.tmp_dir = Path(Env.get().create_tmp_path())
+        self.media_files = {}
+
+    def restore(self, bundle_path: Path) -> 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:
+            # 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)
+
+        pxt_md_version = self.md['pxt_md_version']
+        assert isinstance(pxt_md_version, int)
+
+        if pxt_md_version != metadata.VERSION:
+            raise excs.Error(
+                f'Pixeltable metadata version mismatch: {pxt_md_version} != {metadata.VERSION}.\n'
+                'Please upgrade Pixeltable to use this dataset: pip install -U pixeltable'
+            )
+
+        tbl_md = [schema.FullTableMd.from_dict(t) for t in self.md['md']['tables']]
+
+        # 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
+
+        # 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
+
+        # 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)
+
+        return replica_tbl
+
+    def __import_table(self, bundle_path: Path, tv: catalog.TableVersion, tbl_md: schema.FullTableMd) -> None:
+        """
+        Import the Parquet table into the Pixeltable catalog.
+        """
+        tbl_id = uuid.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():
+            pydict = batch.to_pydict()
+            rows = self.__from_pa_pydict(tv, pydict)
+            tv.store_tbl.load_rows(rows)
+
+    def __from_pa_pydict(self, tv: catalog.TableVersion, pydict: dict[str, Any]) -> list[dict[str, Any]]:
+        # Data conversions from pyarrow to Pixeltable
+        sql_types: dict[str, sql.types.TypeEngine[Any]] = {}
+        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
+
+        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)
+
+        return rows
+
+    def __from_pa_value(
+        self, tv: catalog.TableVersion, val: Any, sql_type: sql.types.TypeEngine[Any], media_col_id: Optional[int]
+    ) -> Any:
+        if val is None:
+            return None
+        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)
+        return val
+
+    def __relocate_media_file(self, tv: catalog.TableVersion, media_col_id: int, url: str) -> str:
+        # If this is a pxtmedia:// URL, relocate it
+        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':
+            if url not in self.media_files:
+                # 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)))
+            return self.media_files[url]
+        # For any type of URL other than a local file, just return the URL as-is.
+        return url

pixeltable/share/publish.py

@@ -9,10 +9,9 @@ from tqdm import tqdm
 import pixeltable as pxt
 from pixeltable import exceptions as excs
 from pixeltable.env import Env
-from pixeltable.metadata.schema import FullTableMd
 from pixeltable.utils import sha256sum
 
-from .packager import TablePackager
+from .packager import TablePackager, TableRestorer
 
 # These URLs are abstracted out for now, but will be replaced with actual (hard-coded) URLs once the
 # pixeltable.com URLs are available.
@@ -20,7 +19,10 @@ from .packager import TablePackager
 PIXELTABLE_API_URL = 'https://internal-api.pixeltable.com'
 
 
-def publish_snapshot(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
+def push_replica(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
+    if not src_tbl._tbl_version.get().is_snapshot:
+        raise excs.Error('Only snapshots may be published.')
+
     packager = TablePackager(src_tbl, additional_md={'table_uri': dest_tbl_uri})
     request_json = packager.md | {'operation_type': 'publish_snapshot'}
     headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
@@ -65,18 +67,6 @@ def publish_snapshot(dest_tbl_uri: str, src_tbl: pxt.Table) -> str:
     return confirmed_tbl_uri
 
 
-def clone_snapshot(dest_tbl_uri: str) -> list[FullTableMd]:
-    headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
-    clone_request_json = {'operation_type': 'clone_snapshot', 'table_uri': dest_tbl_uri}
-    response = requests.post(PIXELTABLE_API_URL, json=clone_request_json, headers=headers_json)
-    if response.status_code != 200:
-        raise excs.Error(f'Error cloning snapshot: {response.text}')
-    response_json = response.json()
-    if not isinstance(response_json, dict) or 'table_uri' not in response_json:
-        raise excs.Error(f'Unexpected response from server.\n{response_json}')
-    return [FullTableMd.from_dict(t) for t in response_json['md']['tables']]
-
-
 def _upload_bundle_to_s3(bundle: Path, parsed_location: urllib.parse.ParseResult) -> None:
     from pixeltable.utils.s3 import get_client
 
@@ -102,5 +92,60 @@ def _upload_bundle_to_s3(bundle: Path, parsed_location: urllib.parse.ParseResult
         file=sys.stdout,
     )
     s3_client.upload_file(
-        Filename=str(bundle), Bucket=bucket, Key=str(remote_path), ExtraArgs=upload_args, Callback=progress_bar.update
+        Filename=str(bundle), Bucket=bucket, Key=remote_path, ExtraArgs=upload_args, Callback=progress_bar.update
+    )
+
+
+def pull_replica(dest_path: str, src_tbl_uri: str) -> pxt.Table:
+    headers_json = {'X-api-key': Env.get().pxt_api_key, 'Content-Type': 'application/json'}
+    clone_request_json = {'operation_type': 'clone_snapshot', 'table_uri': src_tbl_uri}
+    response = requests.post(PIXELTABLE_API_URL, json=clone_request_json, headers=headers_json)
+    if response.status_code != 200:
+        raise excs.Error(f'Error cloning snapshot: {response.text}')
+    response_json = response.json()
+    if not isinstance(response_json, dict) or 'table_uri' not in response_json:
+        raise excs.Error(f'Error cloning shapshot: unexpected response from server.\n{response_json}')
+
+    primary_tbl_additional_md = response_json['md']['tables'][0]['table_md']['additional_md']
+    bundle_uri = primary_tbl_additional_md['destination_uri']
+    bundle_filename = primary_tbl_additional_md['datafile']
+    parsed_location = urllib.parse.urlparse(bundle_uri)
+    if parsed_location.scheme == 's3':
+        bundle_path = _download_bundle_from_s3(parsed_location, bundle_filename)
+    else:
+        raise excs.Error(f'Unexpected response from server: unsupported bundle uri: {bundle_uri}')
+
+    restorer = TableRestorer(dest_path, response_json)
+    tbl = restorer.restore(bundle_path)
+    Env.get().console_logger.info(f'Created local replica {tbl._path!r} from URI: {src_tbl_uri}')
+    return tbl
+
+
+def _download_bundle_from_s3(parsed_location: urllib.parse.ParseResult, bundle_filename: str) -> Path:
+    from pixeltable.utils.s3 import get_client
+
+    bucket = parsed_location.netloc
+    remote_dir = Path(urllib.parse.unquote(urllib.request.url2pathname(parsed_location.path)))
+    remote_path = str(remote_dir / bundle_filename)[1:]  # Remove initial /
+
+    Env.get().console_logger.info(f'Downloading snapshot from: {bucket}:{remote_path}')
+
+    boto_config = {'max_pool_connections': 5, 'connect_timeout': 15, 'retries': {'max_attempts': 3, 'mode': 'adaptive'}}
+    s3_client = get_client(**boto_config)
+
+    obj = s3_client.head_object(Bucket=bucket, Key=remote_path)  # Check if the object exists
+    bundle_size = obj['ContentLength']
+
+    bundle_path = Path(Env.get().create_tmp_path())
+    progress_bar = tqdm(
+        desc='Downloading',
+        total=bundle_size,
+        unit='B',
+        unit_scale=True,
+        unit_divisor=1024,
+        miniters=1,
+        ncols=100,
+        file=sys.stdout,
     )
+    s3_client.download_file(Bucket=bucket, Key=remote_path, Filename=str(bundle_path), Callback=progress_bar.update)
+    return bundle_path

pixeltable/store.py

@@ -7,8 +7,9 @@ import sys
 import urllib.parse
 import urllib.request
 import warnings
-from typing import Any, Iterator, Literal, Optional, Union
+from typing import Any, Iterable, Iterator, Literal, Optional, Union
 
+import more_itertools
 import sqlalchemy as sql
 from tqdm import TqdmWarning, tqdm
 
@@ -190,34 +191,29 @@ class StoreBase:
         assert col.is_stored
         conn = Env.get().conn
         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)
+        s_txt = f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.store_name()} {col_type_str} NULL'
         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()} ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL'
-            )
-            conn.execute(stmt)
-            stmt = sql.text(
-                f'ALTER TABLE {self._storage_name()} ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL'
-            )
-            conn.execute(stmt)
+            s_txt += f' , ADD COLUMN {col.errormsg_store_name()} VARCHAR DEFAULT NULL'
+            s_txt += f' , ADD COLUMN {col.errortype_store_name()} VARCHAR DEFAULT NULL'
             added_storage_cols.extend([col.errormsg_store_name(), col.errortype_store_name()])
+
+        stmt = sql.text(s_txt)
+        log_stmt(_logger, stmt)
+        conn.execute(stmt)
         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) -> None:
         """Execute Alter Table Drop Column statement"""
-        conn = Env.get().conn
-        stmt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.store_name()}'
-        conn.execute(sql.text(stmt))
+        s_txt = f'ALTER TABLE {self._storage_name()} DROP COLUMN {col.store_name()}'
         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))
+            s_txt += f' , DROP COLUMN {col.errormsg_store_name()}'
+            s_txt += f' , DROP COLUMN {col.errortype_store_name()}'
+        stmt = sql.text(s_txt)
+        log_stmt(_logger, stmt)
+        Env.get().conn.execute(stmt)
 
     def load_column(
         self, col: catalog.Column, exec_plan: ExecNode, value_expr_slot_idx: int, on_error: Literal['abort', 'ignore']
@@ -313,7 +309,7 @@ class StoreBase:
     def insert_rows(
         self,
         exec_plan: ExecNode,
-        v_min: Optional[int] = None,
+        v_min: int,
         show_progress: bool = True,
         rowids: Optional[Iterator[int]] = None,
         abort_on_exc: bool = False,
@@ -432,6 +428,35 @@ class StoreBase:
         status = conn.execute(stmt)
         return status.rowcount
 
+    def dump_rows(self, version: int, filter_view: StoreBase, filter_view_version: int) -> Iterator[dict[str, Any]]:
+        filter_predicate = sql.and_(
+            filter_view.v_min_col <= filter_view_version,
+            filter_view.v_max_col > filter_view_version,
+            *[c1 == c2 for c1, c2 in zip(self.rowid_columns(), filter_view.rowid_columns())],
+        )
+        stmt = (
+            sql.select('*')  # TODO: Use a more specific list of columns?
+            .select_from(self.sa_tbl)
+            .where(self.v_min_col <= version)
+            .where(self.v_max_col > version)
+            .where(sql.exists().where(filter_predicate))
+        )
+        conn = Env.get().conn
+        _logger.debug(stmt)
+        log_explain(_logger, stmt, conn)
+        result = conn.execute(stmt)
+        for row in result:
+            yield dict(zip(result.keys(), row))
+
+    def load_rows(self, rows: Iterable[dict[str, Any]], batch_size: int = 10_000) -> None:
+        """
+        When instantiating a replica, we can't rely on the usual insertion code path, which contains error handling
+        and other logic that doesn't apply.
+        """
+        conn = Env.get().conn
+        for batch in more_itertools.batched(rows, batch_size):
+            conn.execute(sql.insert(self.sa_tbl), batch)
+
 
 class StoreTable(StoreBase):
     def __init__(self, tbl_version: catalog.TableVersion):