Repository: xataio/pgstream
Pull request #403 (2025-06-12 10:44:18)

[Snapshot] Automatic configuration of batch pages
This PR updates the snapshot batch configuration to rely on the batch size in bytes instead of the number of pages per batch. The reason for this is that different tables will have variable row sizes, which means the page sizes will be different as well. If the same number of pages is selected for all tables, it will end up being a setting that is not ideal for all tables. 

For example, if there's a table with an average page size of 10KB and another table with an average page size of 10MB, a setting of 100 batch page size would mean for the first table we'll be retrieving  1MB of data per batch, while for the second table that would be 1GB, leading to very different performances and memory usage for a single setting. 

By using the batch bytes as a configuration setting instead, the number of pages per batch will be automatically computed to match that value, producing different values for each table. To compute the number of pages we use the number of pages for a table, and the average page bytes.


Patch:
diff --git a/cmd/config/config_env.go b/cmd/config/config_env.go
index 2ad7600..f1d1330 100644
--- a/cmd/config/config_env.go
+++ b/cmd/config/config_env.go
@@ -42,7 +42,7 @@ func init() {
 	viper.BindEnv("PGSTREAM_POSTGRES_REPLICATION_SLOT_NAME")
 
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL")
-	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE")
+	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS")
 	viper.BindEnv("PGSTREAM_POSTGRES_SNAPSHOT_TABLES")
@@ -197,7 +197,7 @@ func parseSnapshotConfig(pgURL string) *snapshotbuilder.SnapshotListenerConfig {
 	cfg := &snapshotbuilder.SnapshotListenerConfig{
 		Generator: pgsnapshotgenerator.Config{
 			URL:             pgURL,
-			BatchPageSize:   viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE"),
+			BatchBytes:      viper.GetUint64("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES"),
 			SchemaWorkers:   viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS"),
 			TableWorkers:    viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS"),
 			SnapshotWorkers: viper.GetUint("PGSTREAM_POSTGRES_SNAPSHOT_WORKERS"),
diff --git a/cmd/config/config_yaml.go b/cmd/config/config_yaml.go
index d1f890e..fb60d8c 100644
--- a/cmd/config/config_yaml.go
+++ b/cmd/config/config_yaml.go
@@ -86,9 +86,9 @@ type SnapshotRecorderConfig struct {
 }
 
 type SnapshotDataConfig struct {
-	SchemaWorkers int `mapstructure:"schema_workers" yaml:"schema_workers"`
-	TableWorkers  int `mapstructure:"table_workers" yaml:"table_workers"`
-	BatchPageSize int `mapstructure:"batch_page_size" yaml:"batch_page_size"`
+	SchemaWorkers int    `mapstructure:"schema_workers" yaml:"schema_workers"`
+	TableWorkers  int    `mapstructure:"table_workers" yaml:"table_workers"`
+	BatchBytes    uint64 `mapstructure:"batch_bytes" yaml:"batch_bytes"`
 }
 
 type SnapshotSchemaConfig struct {
@@ -456,7 +456,7 @@ func (c *YAMLConfig) parseSnapshotConfig() (*snapshotbuilder.SnapshotListenerCon
 	if snapshotConfig.Mode == fullSnapshotMode || snapshotConfig.Mode == dataSnapshotMode {
 		streamCfg.Generator = pgsnapshotgenerator.Config{
 			URL:             c.Source.Postgres.URL,
-			BatchPageSize:   uint(snapshotConfig.Data.BatchPageSize),
+			BatchBytes:      snapshotConfig.Data.BatchBytes,
 			SchemaWorkers:   uint(snapshotConfig.Data.SchemaWorkers),
 			TableWorkers:    uint(snapshotConfig.Data.TableWorkers),
 			SnapshotWorkers: uint(snapshotConfig.SnapshotWorkers),
diff --git a/config_template.yaml b/config_template.yaml
index 492c51c..d959f68 100644
--- a/config_template.yaml
+++ b/config_template.yaml
@@ -12,7 +12,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
diff --git a/docs/README.md b/docs/README.md
index 2527f3e..7a21fa6 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -107,7 +107,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel. Defaults to 4
         table_workers: 4 # number of workers to snapshot a table in parallel. Defaults to 4
-        batch_page_size: 1000 # number of pages to read per batch. Defaults to 1000
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
@@ -239,7 +239,7 @@ Here's a list of all the environment variables that can be used to configure the
 | PGSTREAM_POSTGRES_SNAPSHOT_TABLES           | ""                           | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
 | PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS   | 4                            | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
 | PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS    | 4                            | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
-| PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE  | 1000                         | No       | Size of the table page range which will be processed concurrently by the table workers from `PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS`.                                                                                                                                                                      |
+| PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES      | 83886080 (80MiB)             | No       | Max batch size in bytes to be read and processed by each table worker at a time. The number of pages in the select queries will be based on this value.                                                                                                                                                      |
 | PGSTREAM_POSTGRES_SNAPSHOT_WORKERS          | 1                            | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
 | PGSTREAM_POSTGRES_SNAPSHOT_USE_SCHEMALOG    | False                        | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres targets.                                                                                                                                                                            |
 | PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL        | ""                           | No       | Postgres URL for the database where the snapshot requests and their status will be tracked. A table `snapshot_requests` will be created under a `pgstream` schema.                                                                                                                                           |
@@ -250,19 +250,19 @@ Here's a list of all the environment variables that can be used to configure the
 <details>
   <summary>Postgres Snapshoter</summary>
 
-| Environment Variable                                 | Default | Required | Description                                                                                                                                                                                                                                                                                                  |
-| ---------------------------------------------------- | ------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL              | N/A     | Yes      | URL of the Postgres database to snapshot.                                                                                                                                                                                                                                                                    |
-| PGSTREAM_POSTGRES_SNAPSHOT_TABLES                    | ""      | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
-| PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS            | 4       | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
-| PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS             | 4       | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
-| PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE           | 1000    | No       | Size of the table page range which will be processed concurrently by the table workers from `PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS`.                                                                                                                                                                      |
-| PGSTREAM_POSTGRES_SNAPSHOT_WORKERS                   | 1       | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
-| PGSTREAM_POSTGRES_SNAPSHOT_USE_SCHEMALOG             | False   | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres targets.                                                                                                                                                                            |
-| PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB           | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to issue commands to DROP all the objects that will be restored.                                                                                                                                                           |
-| PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to snapshot all global database objects outside of the selected schema (such as extensions, triggers, etc).                                                                                                                |
-| PGSTREAM_POSTGRES_SNAPSHOT_CREATE_TARGET_DB          | False   | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to create the database being restored.                                                                                                                                                                                     |
-| PGSTREAM_POSTGRES_SNAPSHOT_ROLE                      | ""      | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, role name to be used to create the dump.                                                                                                                                                                                          |
+| Environment Variable                                 | Default          | Required | Description                                                                                                                                                                                                                                                                                                  |
+| ---------------------------------------------------- | ---------------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL              | N/A              | Yes      | URL of the Postgres database to snapshot.                                                                                                                                                                                                                                                                    |
+| PGSTREAM_POSTGRES_SNAPSHOT_TABLES                    | ""               | No       | Tables for which there will be an initial snapshot generated. The syntax supports wildcards. Tables without a schema defined will be applied the public schema. Example: for `public.test_table` and all tables in the `test_schema` schema, the value would be the following: `"test_table test_schema.\*"` |
+| PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS            | 4                | No       | Number of tables per schema that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                  |
+| PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS             | 4                | No       | Number of concurrent workers that will be used per table by the snapshotting process.                                                                                                                                                                                                                        |
+| PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES               | 83886080 (80MiB) | No       | Max batch size in bytes to be read and processed by each table worker at a time. The number of pages in the select queries will be based on this value.                                                                                                                                                      |
+| PGSTREAM_POSTGRES_SNAPSHOT_WORKERS                   | 1                | No       | Number of schemas that will be processed in parallel by the snapshotting process.                                                                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_SNAPSHOT_USE_SCHEMALOG             | False            | No       | Forces the use of the `pgstream.schema_log` for the schema snapshot instead of using `pg_dump`/`pg_restore` for Postgres targets.                                                                                                                                                                            |
+| PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB           | False            | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to issue commands to DROP all the objects that will be restored.                                                                                                                                                           |
+| PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS | False            | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to snapshot all global database objects outside of the selected schema (such as extensions, triggers, etc).                                                                                                                |
+| PGSTREAM_POSTGRES_SNAPSHOT_CREATE_TARGET_DB          | False            | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, option to create the database being restored.                                                                                                                                                                                     |
+| PGSTREAM_POSTGRES_SNAPSHOT_ROLE                      | ""               | No       | When using `pg_dump`/`pg_restore` to snapshot schema for Postgres targets, role name to be used to create the dump.                                                                                                                                                                                          |
 
 </details>
 
@@ -1118,22 +1118,22 @@ transformations:
 | -------------------------- |
 | json, jsonb                |
 
-| Parameter   | Type   | Default | Required |
-| ----------- | ------ | ------- | -------- |
-| operations  | array  | N/A     | Yes      |
+| Parameter  | Type  | Default | Required |
+| ---------- | ----- | ------- | -------- |
+| operations | array | N/A     | Yes      |
 
 Parameter for each operation:
-| Parameter       | Type    | Default | Required | Values                         |
+| Parameter | Type | Default | Required | Values |
 | --------------- | ------- | ------- | -------- | ------------------------------ |
-| operation       | string  | N/A     | Yes      | set, delete                    |
-| path            | string  | N/A     | Yes      | *sjson syntax                  |
-| error_not_exist | boolean | false   | No       | true, false                    |
-| value           | string  | N\A     | **Yes    | Any valid JSON representation  |
-| value_template  | string  | N\A     | **Yes    | Any template with valid syntax |
+| operation | string | N/A | Yes | set, delete |
+| path | string | N/A | Yes | \*sjson syntax |
+| error_not_exist | boolean | false | No | true, false |
+| value | string | N\A | **Yes | Any valid JSON representation |
+| value_template | string | N\A | **Yes | Any template with valid syntax |
 
-*Paths should follow [sjson syntax](https://github.com/tidwall/sjson#path-syntax)
+\*Paths should follow [sjson syntax](https://github.com/tidwall/sjson#path-syntax)
 
-**Either `value` or `value_template` must be provided if the operation is `set`. If both are provided, `value_template` takes precedence.
+\*\*Either `value` or `value_template` must be provided if the operation is `set`. If both are provided, `value_template` takes precedence.
 
 JSON transformer can be used for Postgres types json and jsonb. This transformer executes a list of given operations on the json data to be transformed.
 
@@ -1145,6 +1145,7 @@ Execution of an operation errors out if the given path does not exists and the p
 JSON transformer uses [sjson](https://github.com/tidwall/sjson) library for executing the operations. Operation paths should follow the [synxtax rules of sjson](https://github.com/tidwall/sjson#path-syntax)
 
 With the below config `pgstream` transforms the json values in the column `user_info_json` of the table `users` by:
+
 - First, traversing all the items in the array named `purchases`, and for each element, setting value to "-" for key "item".
 - Then, deleting the object named "country" under the top-level object "address".
 - Completely masking the "city" value under object "address", using `go-masker`'s default masking function supported by `pgstream`'s templating.
@@ -1172,13 +1173,14 @@ transformations:
                 path: "address.country"
               - operation: set
                 path: "address.city"
-                value_template:  "\"{{ masking \"default\" .GetValue }}\""
+                value_template: '"{{ masking "default" .GetValue }}"'
               - operation: set
                 path: "user.lastname"
-                value_template:  "\"{{ .GetDynamicValue \"lastname\" }}\""
+                value_template: '"{{ .GetDynamicValue "lastname" }}"'
 ```
 
 For input JSON value,
+
 ```json
 {
   "user": {
@@ -1203,6 +1205,7 @@ For input JSON value,
 ```
 
 the JSON transformer with above config produces output:
+
 ```json
 {
   "user": {
@@ -1227,7 +1230,6 @@ the JSON transformer with above config produces output:
 
 </details>
 
-
  <details>
   <summary>literal_string</summary>
 
diff --git a/docs/tutorials/postgres_snapshot.md b/docs/tutorials/postgres_snapshot.md
index 5db01ce..5069a42 100644
--- a/docs/tutorials/postgres_snapshot.md
+++ b/docs/tutorials/postgres_snapshot.md
@@ -71,8 +71,8 @@ Further configuration can be provided to optimize the performance of the snapsho
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 # Number of concurrent workers that will be used per table by the snapshotting process. Defaults to 4.
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-# Size of the table page range which will be processed concurrently by the table workers from PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS. Defaults to 1000.
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+# Size of the batches the table workers from PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS will read concurrently. Defaults to 80MiB.
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080
 # Number of schemas that will be processed in parallel by the snapshotting process. Defaults to 1.
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 ```
@@ -124,7 +124,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL="postgres://postgres:postgres@localhost:
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost:5432?sslmode=disable"
 
@@ -151,7 +151,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
diff --git a/docs/tutorials/snapshot2pg_tutorial.env b/docs/tutorials/snapshot2pg_tutorial.env
index 292e197..1f6314f 100644
--- a/docs/tutorials/snapshot2pg_tutorial.env
+++ b/docs/tutorials/snapshot2pg_tutorial.env
@@ -3,7 +3,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL="postgres://postgres:postgres@localhost:
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost:5432?sslmode=disable"
 
diff --git a/docs/tutorials/snapshot2pg_tutorial.yaml b/docs/tutorials/snapshot2pg_tutorial.yaml
index 3d14678..0cbb1af 100644
--- a/docs/tutorials/snapshot2pg_tutorial.yaml
+++ b/docs/tutorials/snapshot2pg_tutorial.yaml
@@ -12,7 +12,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
diff --git a/pg2kafka.env b/pg2kafka.env
index 898f733..338444e 100644
--- a/pg2kafka.env
+++ b/pg2kafka.env
@@ -3,7 +3,7 @@ PGSTREAM_POSTGRES_LISTENER_URL="postgres://postgres:postgres@localhost?sslmode=d
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080 # 80MiB
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 
 # Processor config
diff --git a/pg2kafka.yaml b/pg2kafka.yaml
index 372d6bd..fa15dbf 100644
--- a/pg2kafka.yaml
+++ b/pg2kafka.yaml
@@ -11,7 +11,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: schemalog # options are pgdump_pgrestore or schemalog
 
diff --git a/pg2os.env b/pg2os.env
index f53014f..dc1e596 100644
--- a/pg2os.env
+++ b/pg2os.env
@@ -3,7 +3,7 @@ PGSTREAM_POSTGRES_LISTENER_URL="postgres://postgres:postgres@localhost?sslmode=d
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080 # 80MiB
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 
 # Processor config
diff --git a/pg2os.yaml b/pg2os.yaml
index 696d1c1..f6d4198 100644
--- a/pg2os.yaml
+++ b/pg2os.yaml
@@ -11,7 +11,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: schemalog # options are pgdump_pgrestore or schemalog
 
diff --git a/pg2pg.env b/pg2pg.env
index a011c59..d043a6f 100644
--- a/pg2pg.env
+++ b/pg2pg.env
@@ -4,7 +4,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_STORE_URL="postgres://postgres:postgres@localhost?ssl
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="*"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080 # 80MiB
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 
 # Processor config
diff --git a/pg2pg.yaml b/pg2pg.yaml
index 8d282d7..5fd3bda 100644
--- a/pg2pg.yaml
+++ b/pg2pg.yaml
@@ -11,7 +11,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
diff --git a/pkg/snapshot/generator/postgres/data/config.go b/pkg/snapshot/generator/postgres/data/config.go
index 9bb797e..faffa6d 100644
--- a/pkg/snapshot/generator/postgres/data/config.go
+++ b/pkg/snapshot/generator/postgres/data/config.go
@@ -5,9 +5,9 @@ package postgres
 type Config struct {
 	// Postgres connection URL. Required.
 	URL string
-	// BatchPageSize represents the size of the table page range that will be
-	// processed concurrently by the table workers. Defaults to 1000.
-	BatchPageSize uint
+	// BatchBytes represents the size of the batch of table pages in bytes.
+	// Defaults to 80MiB.
+	BatchBytes uint64
 	// SnapshotWorkers represents the number of snapshots the generator will
 	// process concurrently. This doesn't affect the parallelism of the tables
 	// within each individual snapshot request. It defaults to 1.
@@ -21,17 +21,17 @@ type Config struct {
 }
 
 const (
-	defaultBatchPageSize   = 1000
 	defaultTableWorkers    = 4
 	defaultSchemaWorkers   = 4
 	defaultSnapshotWorkers = 1
+	defaultBatchBytes      = 80 * 1024 * 1024 // 80 MiB
 )
 
-func (c *Config) batchPageSize() uint {
-	if c.BatchPageSize > 0 {
-		return c.BatchPageSize
+func (c *Config) batchBytes() uint64 {
+	if c.BatchBytes > 0 {
+		return c.BatchBytes
 	}
-	return defaultBatchPageSize
+	return defaultBatchBytes
 }
 
 func (c *Config) schemaWorkers() uint {
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
index 90397ea..151cd5b 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator.go
@@ -27,8 +27,8 @@ type SnapshotGenerator struct {
 	// workers per schema, parallelise the snapshot creation for each table
 	schemaWorkers uint
 	// workers per table, parallelise the snapshot creation for each page range
-	tableWorkers  uint
-	batchPageSize uint
+	tableWorkers uint
+	batchBytes   uint64
 
 	// Function called for processing produced rows.
 	rowsProcessor          snapshot.RowsProcessor
@@ -39,6 +39,11 @@ type mapper interface {
 	TypeForOID(context.Context, uint32) (string, error)
 }
 
+type tablePageInfo struct {
+	pageCount    int
+	avgPageBytes int64
+}
+
 type pageRange struct {
 	start uint
 	end   uint
@@ -64,7 +69,7 @@ func NewSnapshotGenerator(ctx context.Context, cfg *Config, rowsProcessor snapsh
 		mapper:          pglib.NewMapper(conn),
 		conn:            conn,
 		rowsProcessor:   rowsProcessor,
-		batchPageSize:   cfg.batchPageSize(),
+		batchBytes:      cfg.batchBytes(),
 		tableWorkers:    cfg.tableWorkers(),
 		schemaWorkers:   cfg.schemaWorkers(),
 		snapshotWorkers: cfg.snapshotWorkers(),
@@ -236,12 +241,12 @@ func (sg *SnapshotGenerator) collectSchemaErrors(workerSchemaErrs map[string]err
 }
 
 func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID string, schema, table string) error {
-	tablePageCount, err := sg.getTablePageCount(ctx, schema, table, snapshotID)
+	tablePageInfo, err := sg.getTablePageInfo(ctx, schema, table, snapshotID)
 	if err != nil {
 		return err
 	}
 	// the table is empty
-	if tablePageCount < 0 {
+	if tablePageInfo.pageCount < 0 {
 		return nil
 	}
 
@@ -249,7 +254,7 @@ func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID strin
 	// snapshot relies on the transaction snapshot id to ensure all workers
 	// have the same table view, which allows us to use the ctid to
 	// parallelise the work.
-	rangeChan := make(chan pageRange, tablePageCount)
+	rangeChan := make(chan pageRange, tablePageInfo.pageCount)
 	errGroup, ctx := errgroup.WithContext(ctx)
 	for i := uint(0); i < sg.tableWorkers; i++ {
 		errGroup.Go(func() error {
@@ -257,12 +262,19 @@ func (sg *SnapshotGenerator) snapshotTable(ctx context.Context, snapshotID strin
 		})
 	}
 
+	// automatically determine the batch page size based on the average page
+	// size and the configured batch bytes limit.
+	batchPageSize := tablePageInfo.getBatchPageSize(sg.batchBytes)
+	sg.logger.Debug("batch page size for table", loglib.Fields{
+		"schema": schema, "table": table, "batch_page_size": batchPageSize,
+		"page_count": tablePageInfo.pageCount, "avg_page_bytes": tablePageInfo.avgPageBytes,
+	})
 	// page count returned by postgres starts at 0, so we need to include it
 	// when creating the page ranges.
-	for start := uint(0); start <= uint(tablePageCount); start += sg.batchPageSize {
+	for start := uint(0); start <= uint(tablePageInfo.pageCount); start += batchPageSize {
 		rangeChan <- pageRange{
 			start: start,
-			end:   start + sg.batchPageSize,
+			end:   start + batchPageSize,
 		}
 	}
 
@@ -280,14 +292,15 @@ func (sg *SnapshotGenerator) snapshotTableRangeWorker(ctx context.Context, snaps
 	return nil
 }
 
+var pageRangeQuery = "SELECT * FROM %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'"
+
 func (sg *SnapshotGenerator) snapshotTableRange(ctx context.Context, snapshotID, schema, table string, pageRange pageRange) error {
 	return sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
 		sg.logger.Debug(fmt.Sprintf("querying table page range %d-%d", pageRange.start, pageRange.end), loglib.Fields{
 			"schema": schema, "table": table, "snapshotID": snapshotID,
 		})
 
-		query := fmt.Sprintf("SELECT * FROM %s WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'",
-			pglib.QuoteQualifiedIdentifier(schema, table), pageRange.start, pageRange.end)
+		query := fmt.Sprintf(pageRangeQuery, pglib.QuoteQualifiedIdentifier(schema, table), pageRange.start, pageRange.end)
 		rows, err := tx.Query(ctx, query)
 		if err != nil {
 			return fmt.Errorf("querying table rows: %w", err)
@@ -347,29 +360,43 @@ func (sg *SnapshotGenerator) toSnapshotColumns(ctx context.Context, fieldDescrip
 	return columns
 }
 
-func (sg *SnapshotGenerator) getTablePageCount(ctx context.Context, schemaName, tableName, snapshotID string) (int, error) {
-	pageCount := int(0)
+// use pg_table_size instead of pg_total_relation_size since we only care about the size of the table itself and toast tables, not indices.
+// pg_relation_size will return only the size of the table itself, without toast tables.
+const tablePageInfoQuery = `SELECT
+  c.relpages AS page_count,
+  (pg_table_size(c.oid) / COALESCE(NULLIF(c.relpages, 0),1)) AS avg_page_size_bytes
+FROM
+  pg_class c
+  JOIN pg_namespace n ON n.oid = c.relnamespace
+WHERE
+  c.relname = $1
+  AND n.nspname = $2
+  AND c.relkind = 'r';`
+
+func (sg *SnapshotGenerator) getTablePageInfo(ctx context.Context, schemaName, tableName, snapshotID string) (*tablePageInfo, error) {
+	tablePageInfo := &tablePageInfo{}
 	err := sg.execInSnapshotTx(ctx, snapshotID, func(tx pglib.Tx) error {
-		const query = "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2"
-		if err := tx.QueryRow(ctx, query, tableName, schemaName).Scan(&pageCount); err != nil {
-			return fmt.Errorf("getting page count for table %s.%s: %w", schemaName, tableName, err)
+		if err := tx.QueryRow(ctx, tablePageInfoQuery, tableName, schemaName).Scan(&tablePageInfo.pageCount, &tablePageInfo.avgPageBytes); err != nil {
+			return fmt.Errorf("getting page information for table %s.%s: %w", schemaName, tableName, err)
 		}
 
-		sg.logger.Debug(fmt.Sprintf("table page count: %d", pageCount), loglib.Fields{
+		sg.logger.Debug(fmt.Sprintf("table page count: %d, page avg bytes: %d", tablePageInfo.pageCount, tablePageInfo.avgPageBytes), loglib.Fields{
 			"schema": schemaName, "table": tableName, "snapshotID": snapshotID,
 		})
 		return nil
 	})
 	if err != nil {
-		return 0, err
+		return nil, err
 	}
 
-	return pageCount, nil
+	return tablePageInfo, nil
 }
 
+const exportSnapshotQuery = `SELECT pg_export_snapshot()`
+
 func (sg *SnapshotGenerator) exportSnapshot(ctx context.Context, tx pglib.Tx) (string, error) {
 	var snapshotID string
-	if err := tx.QueryRow(ctx, "SELECT pg_export_snapshot()").Scan(&snapshotID); err != nil {
+	if err := tx.QueryRow(ctx, exportSnapshotQuery).Scan(&snapshotID); err != nil {
 		return "", fmt.Errorf("exporting snapshot: %w", err)
 	}
 	return snapshotID, nil
@@ -399,3 +426,28 @@ func snapshotTxOptions() pglib.TxOptions {
 		AccessMode:     pglib.ReadOnly,
 	}
 }
+
+func (t *tablePageInfo) getBatchPageSize(bytes uint64) uint {
+	// at least one page is needed to process the table
+	if t.pageCount == 0 {
+		return 1
+	}
+
+	// no limit on bytes, return all pages
+	if bytes == 0 || t.avgPageBytes == 0 {
+		return uint(t.pageCount)
+	}
+
+	batchPageSize := bytes / uint64(t.avgPageBytes)
+	// at least one page is needed to process the table
+	if batchPageSize == 0 {
+		batchPageSize = 1
+	}
+
+	// don't exceed the total page count
+	if batchPageSize > uint64(t.pageCount) {
+		batchPageSize = uint64(t.pageCount)
+	}
+
+	return uint(batchPageSize)
+}
diff --git a/snapshot2pg.env b/snapshot2pg.env
index 76138bf..5ab8484 100644
--- a/snapshot2pg.env
+++ b/snapshot2pg.env
@@ -3,7 +3,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_LISTENER_URL="postgres://postgres:postgres@localhost?
 PGSTREAM_POSTGRES_SNAPSHOT_TABLES="test"
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080 # 80MiB
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=1
 
 # Processor config


Test patch:
diff --git a/cmd/config/config_env_test.go b/cmd/config/config_env_test.go
index 543a0a7..993437f 100644
--- a/cmd/config/config_env_test.go
+++ b/cmd/config/config_env_test.go
@@ -39,7 +39,7 @@ func Test_EnvVarsToStreamConfig(t *testing.T) {
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_WORKERS", "4")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS", "4")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS", "4")
-	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE", "1000")
+	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES", "83886080")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB", "true")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_CREATE_TARGET_DB", "true")
 	os.Setenv("PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS", "true")
diff --git a/cmd/config/helper_test.go b/cmd/config/helper_test.go
index 12bcf9f..75bb7ab 100644
--- a/cmd/config/helper_test.go
+++ b/cmd/config/helper_test.go
@@ -23,7 +23,7 @@ func validateTestStreamConfig(t *testing.T, streamConfig *stream.Config) {
 	assert.Equal(t, uint(4), streamConfig.Listener.Postgres.Snapshot.Generator.SnapshotWorkers)
 	assert.Equal(t, uint(4), streamConfig.Listener.Postgres.Snapshot.Generator.SchemaWorkers)
 	assert.Equal(t, uint(4), streamConfig.Listener.Postgres.Snapshot.Generator.TableWorkers)
-	assert.Equal(t, uint(1000), streamConfig.Listener.Postgres.Snapshot.Generator.BatchPageSize)
+	assert.Equal(t, uint64(83886080), streamConfig.Listener.Postgres.Snapshot.Generator.BatchBytes)
 	assert.NotNil(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore)
 	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore.CleanTargetDB)
 	assert.True(t, streamConfig.Listener.Postgres.Snapshot.Schema.DumpRestore.CreateTargetDB)
diff --git a/cmd/config/test/test_config.env b/cmd/config/test/test_config.env
index 4e96f3c..81c754b 100644
--- a/cmd/config/test/test_config.env
+++ b/cmd/config/test/test_config.env
@@ -9,7 +9,7 @@ PGSTREAM_POSTGRES_SNAPSHOT_TABLES="test test_schema.Test another_schema.*"
 PGSTREAM_POSTGRES_SNAPSHOT_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_SCHEMA_WORKERS=4
 PGSTREAM_POSTGRES_SNAPSHOT_TABLE_WORKERS=4
-PGSTREAM_POSTGRES_SNAPSHOT_BATCH_PAGE_SIZE=1000
+PGSTREAM_POSTGRES_SNAPSHOT_BATCH_BYTES=83886080 # 80MiB
 PGSTREAM_POSTGRES_SNAPSHOT_CLEAN_TARGET_DB=true
 PGSTREAM_POSTGRES_SNAPSHOT_CREATE_TARGET_DB=true
 PGSTREAM_POSTGRES_SNAPSHOT_INCLUDE_GLOBAL_DB_OBJECTS=true
diff --git a/cmd/config/test/test_config.yaml b/cmd/config/test/test_config.yaml
index 862b388..2149cc2 100644
--- a/cmd/config/test/test_config.yaml
+++ b/cmd/config/test/test_config.yaml
@@ -12,7 +12,7 @@ source:
       data: # when mode is full or data
         schema_workers: 4 # number of schema tables to be snapshotted in parallel
         table_workers: 4 # number of workers to snapshot a table in parallel
-        batch_page_size: 1000 # number of pages to read per batch
+        batch_bytes: 83886080 # bytes to read per batch (defaults to 80MiB)
       schema: # when mode is full or schema
         mode: pgdump_pgrestore # options are pgdump_pgrestore or schemalog
         pgdump_pgrestore:
diff --git a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
index aca90ad..b30893f 100644
--- a/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
+++ b/pkg/snapshot/generator/postgres/data/pg_snapshot_generator_test.go
@@ -33,6 +33,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 			testSchema: {testTable1},
 		},
 	}
+	quotedSchemaTable1 := pglib.QuoteQualifiedIdentifier(testSchema, testTable1)
 
 	txOptions := pglib.TxOptions{
 		IsolationLevel: pglib.RepeatableRead,
@@ -40,7 +41,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 	}
 
 	testSnapshotID := "test-snapshot-id"
-	testPageCount := 1
+	testPageCount := 0 // 0 means 1 page
+	testPageAvgBytes := int64(1024)
 	testUUID := uuid.New().String()
 	testColumns := []snapshot.Column{
 		{Name: "id", Type: "uuid", Value: testUUID},
@@ -75,7 +77,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -96,14 +98,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -118,7 +123,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
-								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSchema, testTable1, 0, 10), query)
+								require.Equal(t, fmt.Sprintf(pageRangeQuery, quotedSchemaTable1, 0, 1), query)
 								require.Len(t, args, 0)
 								return &pgmocks.Rows{
 									CloseFn: func() {},
@@ -153,7 +158,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					require.Equal(t, txOptions, to)
 					mockTx := pgmocks.Tx{
 						QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-							if query == "SELECT pg_export_snapshot()" {
+							if query == exportSnapshotQuery {
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -164,13 +169,16 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 									},
 								}
 							}
-							if query == "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2" {
+							if query == tablePageInfoQuery {
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -228,7 +236,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -249,14 +257,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -271,7 +282,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
-								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSchema, testTable1, 0, 10), query)
+								require.Equal(t, fmt.Sprintf(pageRangeQuery, quotedSchemaTable1, 0, 1), query)
 								require.Len(t, args, 0)
 								return &pgmocks.Rows{
 									CloseFn: func() {},
@@ -309,7 +320,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -330,14 +341,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -352,7 +366,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
-								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSchema, testTable1, 0, 10), query)
+								require.Equal(t, fmt.Sprintf(pageRangeQuery, quotedSchemaTable1, 0, 1), query)
 								require.Len(t, args, 0)
 								return &pgmocks.Rows{
 									CloseFn:             func() {},
@@ -382,7 +396,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										return errTest
@@ -409,7 +423,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -456,7 +470,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -477,7 +491,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
@@ -497,7 +511,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				testSchema: &snapshot.SchemaErrors{
 					Schema: testSchema,
 					TableErrors: map[string]string{
-						testTable1: fmt.Sprintf("getting page count for table test-schema.test-table-1: %v", errTest),
+						testTable1: fmt.Sprintf("getting page information for table test-schema.test-table-1: %v", errTest),
 					},
 				},
 			},
@@ -512,7 +526,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -533,14 +547,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -581,7 +598,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -602,14 +619,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -624,7 +644,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryFn: func(ctx context.Context, query string, args ...any) (pglib.Rows, error) {
-								require.Equal(t, fmt.Sprintf("SELECT * FROM %q.%q WHERE ctid BETWEEN '(%d,0)' AND '(%d,0)'", testSchema, testTable1, 0, 10), query)
+								require.Equal(t, fmt.Sprintf(pageRangeQuery, quotedSchemaTable1, 0, 1), query)
 								require.Len(t, args, 0)
 								return nil, errTest
 							},
@@ -655,7 +675,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -676,14 +696,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -733,7 +756,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					case 1:
 						mockTx := pgmocks.Tx{
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT pg_export_snapshot()", query)
+								require.Equal(t, exportSnapshotQuery, query)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -754,14 +777,17 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 								return pglib.CommandTag{}, nil
 							},
 							QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-								require.Equal(t, "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2", query)
+								require.Equal(t, tablePageInfoQuery, query)
 								require.Equal(t, []any{testTable1, testSchema}, args)
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
-										require.Len(t, args, 1)
+										require.Len(t, args, 2)
 										pageCount, ok := args[0].(*int)
 										require.True(t, ok, fmt.Sprintf("pageCount, expected *int, got %T", args[0]))
 										*pageCount = testPageCount
+										pageAvgBytes, ok := args[1].(*int64)
+										require.True(t, ok, fmt.Sprintf("pageAvgBytes, expected *int, got %T", args[1]))
+										*pageAvgBytes = testPageAvgBytes
 										return nil
 									},
 								}
@@ -809,7 +835,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 					require.Equal(t, txOptions, to)
 					mockTx := pgmocks.Tx{
 						QueryRowFn: func(ctx context.Context, query string, args ...any) pglib.Row {
-							if query == "SELECT pg_export_snapshot()" {
+							if query == exportSnapshotQuery {
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										require.Len(t, args, 1)
@@ -820,7 +846,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 									},
 								}
 							}
-							if query == "SELECT c.relpages FROM pg_class c JOIN pg_namespace n ON c.relnamespace=n.oid WHERE c.relname=$1 and n.nspname=$2" {
+							if query == tablePageInfoQuery {
 								return &pgmocks.Row{
 									ScanFn: func(args ...any) error {
 										return errTest
@@ -851,8 +877,8 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				testSchema: &snapshot.SchemaErrors{
 					Schema: testSchema,
 					TableErrors: map[string]string{
-						testTable1: fmt.Sprintf("getting page count for table %s.%s: %v", testSchema, testTable1, errTest),
-						testTable2: fmt.Sprintf("getting page count for table %s.%s: %v", testSchema, testTable2, errTest),
+						testTable1: fmt.Sprintf("getting page information for table %s.%s: %v", testSchema, testTable1, errTest),
+						testTable2: fmt.Sprintf("getting page information for table %s.%s: %v", testSchema, testTable2, errTest),
 					},
 				},
 			},
@@ -879,7 +905,7 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 				},
 				schemaWorkers:   1,
 				tableWorkers:    1,
-				batchPageSize:   10,
+				batchBytes:      1024 * 1024, // 1MB
 				snapshotWorkers: 1,
 			}
 			sg.tableSnapshotGenerator = sg.snapshotTable
@@ -908,3 +934,77 @@ func TestSnapshotGenerator_CreateSnapshot(t *testing.T) {
 		})
 	}
 }
+
+func TestTablePageInfo_getBatchPageSize(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name          string
+		tablePageInfo *tablePageInfo
+		bytes         uint64
+
+		wantPageSize uint
+	}{
+		{
+			name: "no avg page bytes - return total pages",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    1,
+				avgPageBytes: 0,
+			},
+			wantPageSize: 1,
+		},
+		{
+			name: "single page - return total pages",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    0,
+				avgPageBytes: 0,
+			},
+			wantPageSize: 1,
+		},
+		{
+			name: "multiple pages with average page bytes",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    100,
+				avgPageBytes: 10,
+			},
+			bytes:        1000,
+			wantPageSize: 100,
+		},
+		{
+			name: "page size != page count",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    100,
+				avgPageBytes: 10,
+			},
+			bytes:        10,
+			wantPageSize: 1,
+		},
+		{
+			name: "bytes > avgPageBytes",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    100,
+				avgPageBytes: 11,
+			},
+			bytes:        10,
+			wantPageSize: 1,
+		},
+		{
+			name: "batch page size > total pages",
+			tablePageInfo: &tablePageInfo{
+				pageCount:    1,
+				avgPageBytes: 10,
+			},
+			bytes:        1000,
+			wantPageSize: 1,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			pageSize := tc.tablePageInfo.getBatchPageSize(tc.bytes)
+			require.Equal(t, tc.wantPageSize, pageSize, "wanted page size %d, got %d", tc.wantPageSize, pageSize)
+		})
+	}
+}