support ctid partitioning for inherited tables (#4084)

- implement getInheritedTables to recursively discover inherited tables
and their block count, similar to partitioned tables, except that in
this case the root table's blocks are also included
- updated SQL query logic for inherited tables: 
- explicitly list the column names rather than using `*` for when
childTableRanges is not nil (i.e. when it's inherited/partitioned table)
in case child table has additional columns not present in parent. This
is consistent with the behavior today using NTILE.
- use FROM ONLY instead of FROM for querying inherited/partitioned
tables to avoid double-counting for inherited tables; this does not
change behavior for partitioned tables.
- note that for inherited tables, CHildTableRanges include all
partitions (including the root table), so like partitioned tables, the
Range field in QRepPartition will be nil.


Testing:
- Add e2e tests for single-level and multi-level inherited tables, as
well as scenario where child table contains extra columns.
- Run locally e2e (identified a separate bug where CDC ignores change
events from child inherited tables, but this bug existed before this PR)

Author: jgao54

flow/connectors/postgres/qrep.go

@@ -230,16 +230,19 @@ func (c *PostgresConnector) getPartitions(
 	}
 
 	var partitionFunc PartitioningFunc
+	var partitionFuncName string
 	switch {
 	case isCTIDWatermarkCol && (hasCTIDOverride || hasPartitionOverride):
 		partitionFunc = CTIDBlockPartitioningFunc
+		partitionFuncName = "CTIDBlockPartitioningFunc"
 	case hasPartitionOverride:
 		partitionFunc = MinMaxRangePartitioningFunc
+		partitionFuncName = "MinMaxRangePartitioningFunc"
 	default:
 		partitionFunc = NTileBucketPartitioningFunc
+		partitionFuncName = "NTileBucketPartitioningFunc"
 	}
-
-	c.logger.Info("using partition function", slog.String("partitionFunc", fmt.Sprintf("%T", partitionFunc)))
+	c.logger.Info("using partition function", slog.String("partitionFunc", partitionFuncName))
 	return partitionFunc(ctx, partitionParams)
 }
 
@@ -368,7 +371,7 @@ func corePullQRepRecords(
 	partitionIdLog := slog.String(string(shared.PartitionIDKey), partition.PartitionId)
 
 	selectedColumns := "*"
-	if len(config.Exclude) != 0 {
+	if len(config.Exclude) != 0 || len(partition.ChildTableRanges) > 0 {
 		tableSchema, err := internal.LoadTableSchemaFromCatalog(ctx, catalogPool, config.ParentMirrorName, config.DestinationTableIdentifier)
 		if err != nil {
 			return 0, 0, fmt.Errorf("failed to load table schema: %w", err)
@@ -507,7 +510,8 @@ func pullChildTableRanges(
 			return 0, 0, fmt.Errorf("failed to parse child table %s: %w", child.Table, err)
 		}
 
-		query := fmt.Sprintf("SELECT %s FROM %s WHERE %s BETWEEN $1 AND $2",
+		// ONLY excludes rows from child tables, ensuring we don't double-count inherited rows.
+		query := fmt.Sprintf("SELECT %s FROM ONLY %s WHERE %s BETWEEN $1 AND $2",
 			selectedColumns, parsedChild.String(), quotedWatermarkColumn)
 
 		rangeStart := pgtype.TID{

flow/connectors/postgres/qrep_partition.go

@@ -3,6 +3,7 @@ package connpostgres
 import (
 	"cmp"
 	"context"
+	"errors"
 	"fmt"
 	"log/slog"
 	"maps"
@@ -118,54 +119,43 @@ func MinMaxRangePartitioningFunc(ctx context.Context, pp PartitionParams) ([]*pr
 // CTIDBlockPartitioningFunc is a table partition strategy where partitions are created by dividing table
 // blocks uniformly. Note that partition boundaries (block ranges) are uniform, but actual row distribution
 // may be skewed due to table bloat, deleted tuples, or uneven data distribution across blocks.
-// TODO: add support for inherited tables
 func CTIDBlockPartitioningFunc(ctx context.Context, pp PartitionParams) ([]*protos.QRepPartition, error) {
 	if pp.numPartitions <= 0 {
 		return nil, fmt.Errorf("expect numPartitions to be greater than 0")
 	}
 
-	isPartitioned, err := isPartitionedTable(ctx, pp.tx, pp.watermarkTable)
+	tc, err := classifyTable(ctx, pp.tx, pp.watermarkTable)
 	if err != nil {
 		return nil, err
 	}
-	if isPartitioned {
-		blocksPerTable, err := getPartitionedTables(ctx, pp.tx, pp.watermarkTable)
+	switch {
+	case tc.relkind == "p":
+		blocksPerTable, err := getPartitionedTables(ctx, pp.tx, tc.qualifiedName)
 		if err != nil {
-			return nil, fmt.Errorf("failed to get partitioned tables: %w", err)
+			return nil, fmt.Errorf("failed to get child partitioned tables: %w", err)
 		}
-		var totalBlocks int64
-		for _, blocks := range blocksPerTable {
-			totalBlocks += blocks
-		}
-		if totalBlocks == 0 {
-			pp.logger.Warn("all child partitions are empty, returning empty partition list")
-			return nil, nil
+		return ctidPartitionsForChildTables(pp, blocksPerTable)
+	case tc.relhassubclass:
+		blocksPerTable, err := getInheritedTables(ctx, pp.tx, tc.qualifiedName)
+		if err != nil {
+			return nil, fmt.Errorf("failed to get child inherited tables: %w", err)
 		}
-		pp.logger.Info("[CTIDBlockPartitioning] partitioned table detected",
-			slog.String("table", pp.watermarkTable),
-			slog.Int("numPartitionedTable", len(blocksPerTable)),
-			slog.Int64("totalBlocks", totalBlocks))
-		return ctidPartitionsForPartitionedTable(pp, blocksPerTable, totalBlocks)
+		return ctidPartitionsForChildTables(pp, blocksPerTable)
+	default:
+		return ctidPartitionsForTable(ctx, pp)
 	}
+}
 
-	var totalBlocks pgtype.Int8
-	if err := pp.tx.QueryRow(ctx,
-		"SELECT (pg_relation_size(to_regclass($1)) / current_setting('block_size')::int)::bigint",
-		pp.watermarkTable).Scan(&totalBlocks); err != nil {
-		return nil, fmt.Errorf("failed to get relation blocks for %s: %w", pp.watermarkTable, err)
+// ctidPartitionsForTable generates CTID block partitions for a single (non-partitioned) table.
+func ctidPartitionsForTable(ctx context.Context, pp PartitionParams) ([]*protos.QRepPartition, error) {
+	totalBlocks, err := getTableBlockCount(ctx, pp.tx, pp.watermarkTable)
+	if err != nil {
+		return nil, err
 	}
-	if !totalBlocks.Valid || totalBlocks.Int64 == 0 {
+	if totalBlocks == 0 {
 		pp.logger.Warn("table has 0 blocks, returning empty partition list", slog.String("table", pp.watermarkTable))
 		return nil, nil
 	}
-	return ctidPartitionsForTable(pp, totalBlocks.Int64)
-}
-
-// ctidPartitionsForTable generates CTID block partitions for a single (non-partitioned) table.
-func ctidPartitionsForTable(pp PartitionParams, totalBlocks int64) ([]*protos.QRepPartition, error) {
-	if totalBlocks <= 0 {
-		return nil, fmt.Errorf("invalid block count %d for table %s", totalBlocks, pp.watermarkTable)
-	}
 
 	tidCmp := func(a pgtype.TID, b pgtype.TID) int {
 		if blockCmp := cmp.Compare(a.BlockNumber, b.BlockNumber); blockCmp != 0 {
@@ -225,22 +215,36 @@ func ctidPartitionsForTable(pp PartitionParams, totalBlocks int64) ([]*protos.QR
 	return partitionHelper.GetPartitions(), nil
 }
 
-// ctidPartitionsForPartitionedTable generates snapshot partitions for a partitioned table.
-// It walks child tables sequentially, greedily filling each partition with <blocksPerPartition>
-// blocks before starting the next. A single snapshot partition may span multiple child tables,
-// and a single child table may be split across snapshot partitions.
+// ctidPartitionsForChildTables discovers child tables via the provided discovery function,
+// then generates snapshot partitions using a greedy algorithm.
+// It walks tables (sorted by name) sequentially, greedily filling each partition with
+// <blocksPerPartition> blocks before starting the next. A single snapshot partition may
+// span multiple tables, and a single table may be split across snapshot partitions.
 //
-// Example: children [T1:30 blocks, T2:20 blocks, T3:10 blocks] with blocksPerPartition=25 produces:
+// Example: tables [T1:30 blocks, T2:20 blocks, T3:10 blocks] with blocksPerPartition=25 produces:
 //   - partition 1: [T1:0-24]
 //   - partition 2: [T1:25-29, T2:0-19]
 //   - partition 3: [T3:0-9]
-func ctidPartitionsForPartitionedTable(
+func ctidPartitionsForChildTables(
 	pp PartitionParams,
 	blocksPerTable map[string]int64,
-	totalChildBlocks int64,
 ) ([]*protos.QRepPartition, error) {
+	var totalBlocks int64
+	for _, blocks := range blocksPerTable {
+		totalBlocks += blocks
+	}
+	if totalBlocks == 0 {
+		pp.logger.Warn("zero non-empty child tables found")
+		return nil, nil
+	}
+
+	pp.logger.Info("child tables detected",
+		slog.String("table", pp.watermarkTable),
+		slog.Int("numChildTables", len(blocksPerTable)),
+		slog.Int64("totalBlocks", totalBlocks))
+
 	sortedTables := slices.Sorted(maps.Keys(blocksPerTable))
-	blocksPerPartition := max(int64(1), shared.DivCeil(totalChildBlocks, pp.numPartitions))
+	blocksPerPartition := max(int64(1), shared.DivCeil(totalBlocks, pp.numPartitions))
 
 	var partitions []*protos.QRepPartition
 	childIdx := 0
@@ -280,26 +284,47 @@ func ctidPartitionsForPartitionedTable(
 		}
 	}
 
-	pp.logger.Info("[CTIDBlockPartitioning] generated partitions for partitioned table",
+	pp.logger.Info("generated partitions for child tables",
 		slog.Int("numPartitions", len(partitions)),
 		slog.Int64("blocksPerPartition", blocksPerPartition),
-		slog.Int64("totalBlocks", totalChildBlocks))
+		slog.Int64("totalBlocks", totalBlocks))
 
 	return partitions, nil
 }
 
-func isPartitionedTable(ctx context.Context, tx pgx.Tx, table string) (bool, error) {
-	var relkind string
+func getTableBlockCount(ctx context.Context, tx pgx.Tx, table string) (int64, error) {
+	var blocks int64
+	if err := tx.QueryRow(ctx,
+		"SELECT (pg_relation_size(to_regclass($1)) / current_setting('block_size')::int)::bigint",
+		table).Scan(&blocks); err != nil {
+		return 0, fmt.Errorf("failed to get block count for %s: %w", table, err)
+	}
+	return blocks, nil
+}
+
+type tableClassification struct {
+	qualifiedName  string
+	relkind        string
+	relhassubclass bool
+}
+
+func classifyTable(ctx context.Context, tx pgx.Tx, table string) (tableClassification, error) {
+	var classification tableClassification
 	err := tx.QueryRow(ctx,
-		`SELECT c.relkind::text FROM pg_class c WHERE c.oid = to_regclass($1)`,
-		table).Scan(&relkind)
+		// We fetch the qualified name from pg_class instead of using the input directly because
+		// the input table name is quoted; and we want the unquoted canonical form for uniformity
+		`SELECT format('%s.%s', n.nspname, c.relname), c.relkind::text, c.relhassubclass
+		 FROM pg_class c
+		 JOIN pg_namespace n ON c.relnamespace = n.oid
+		 WHERE c.oid = to_regclass($1)`,
+		table).Scan(&classification.qualifiedName, &classification.relkind, &classification.relhassubclass)
 	if err != nil {
-		if err == pgx.ErrNoRows {
-			return false, fmt.Errorf("table %s not found in pg_class", table)
+		if errors.Is(err, pgx.ErrNoRows) {
+			return classification, fmt.Errorf("table %s not found in pg_class", table)
 		}
-		return false, fmt.Errorf("failed to check relkind for %s: %w", table, err)
+		return classification, fmt.Errorf("failed to classify table %s: %w", table, err)
 	}
-	return relkind == "p", nil
+	return classification, nil
 }
 
 // getPartitionedTables returns a map of partitioned child table names to their block counts,
@@ -351,6 +376,66 @@ func getPartitionedTables(ctx context.Context, tx pgx.Tx, table string) (map[str
 	return blocksPerPartitionedTable, nil
 }
 
+// getInheritedTables returns a map of table names to their block counts for an inherited
+// table hierarchy. Unlike partitioned tables, the parent itself stores data and is included.
+func getInheritedTables(ctx context.Context, tx pgx.Tx, table string) (map[string]int64, error) {
+	blocksPerTable := make(map[string]int64)
+
+	numBlocks, err := getTableBlockCount(ctx, tx, table)
+	if err != nil {
+		return nil, err
+	}
+	if numBlocks > 0 {
+		blocksPerTable[table] = numBlocks
+	}
+
+	rows, err := tx.Query(ctx, `
+		SELECT format('%s.%s', n.nspname, c.relname), c.relhassubclass,
+			(pg_relation_size(c.oid) / current_setting('block_size')::int)::bigint
+		FROM pg_inherits i
+		JOIN pg_class c ON i.inhrelid = c.oid
+		JOIN pg_namespace n ON c.relnamespace = n.oid
+		WHERE i.inhparent = to_regclass($1)
+		ORDER BY c.relname
+	`, table)
+	if err != nil {
+		return nil, fmt.Errorf("failed to query child tables for %s: %w", table, err)
+	}
+	defer rows.Close()
+
+	type tableInfo struct {
+		name        string
+		hasSubclass bool
+		blocks      int64
+	}
+	var children []tableInfo
+	for rows.Next() {
+		var info tableInfo
+		if err := rows.Scan(&info.name, &info.hasSubclass, &info.blocks); err != nil {
+			return nil, fmt.Errorf("failed to scan child table: %w", err)
+		}
+		children = append(children, info)
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+
+	for _, child := range children {
+		if child.blocks > 0 {
+			blocksPerTable[child.name] = child.blocks
+		}
+		if child.hasSubclass {
+			childTables, err := getInheritedTables(ctx, tx, child.name)
+			if err != nil {
+				return nil, fmt.Errorf("failed to get inherited tables of %s: %w", child.name, err)
+			}
+			maps.Copy(blocksPerTable, childTables)
+		}
+	}
+
+	return blocksPerTable, nil
+}
+
 // ComputeNumPartitions computes the number of partitions given desired number of rows
 // per partition, with automatic adjustment to respect the maximum partition limit.
 // TODO: use estimated row count for partitioned/inherited tables as well