diff --git a/flow/connectors/mongo/qrep_partition.go b/flow/connectors/mongo/qrep_partition.go index 38c883220..d19a0de7a 100644 --- a/flow/connectors/mongo/qrep_partition.go +++ b/flow/connectors/mongo/qrep_partition.go @@ -139,7 +139,7 @@ func (c *MongoConnector) numericPartitions( return utils.FullTablePartition(), nil } - ranges := computeNumericRanges(minVal, maxVal, numPartitions) + ranges := utils.ComputeRanges(minVal, maxVal, numPartitions) c.logger.Info("[mongo] using numeric _id partitioning", slog.Int("numPartitions", len(ranges))) partitions := make([]*protos.QRepPartition, 0, len(ranges)) for i, rng := range ranges { @@ -148,27 +148,6 @@ func (c *MongoConnector) numericPartitions( return partitions, nil } -func computeNumericRanges(minVal int64, maxVal int64, numPartitions int64) [][2]int64 { - // span arithmetic is done in uint64 to avoid overflow - span := uint64(maxVal) - uint64(minVal) - step := span / uint64(numPartitions) - if span%uint64(numPartitions) != 0 { - step++ - } - - ranges := make([][2]int64, 0, numPartitions) - start := minVal - for { - if remaining := uint64(maxVal) - uint64(start); remaining <= step { - ranges = append(ranges, [2]int64{start, maxVal}) - return ranges - } - end := int64(uint64(start) + step) - ranges = append(ranges, [2]int64{start, end}) - start = end - } -} - // stringPartitions builds partitions for a string _id collection. Because string // keys (e.g. package names) are not uniformly distributed, uniform division of the // keyspace would be badly skewed; instead we sample the collection and pick quantile diff --git a/flow/connectors/mongo/qrep_partition_test.go b/flow/connectors/mongo/qrep_partition_test.go index 408e234d0..e4f3f7468 100644 --- a/flow/connectors/mongo/qrep_partition_test.go +++ b/flow/connectors/mongo/qrep_partition_test.go @@ -1,7 +1,6 @@ package connmongo import ( - "math" "testing" "github.com/stretchr/testify/require" @@ -119,51 +118,6 @@ func TestComputeStringBoundaries(t *testing.T) { }) } -func assertContiguousNumeric(t *testing.T, ranges [][2]int64, minVal, maxVal int64) { - t.Helper() - require.NotEmpty(t, ranges) - require.Equal(t, minVal, ranges[0][0]) - require.Equal(t, maxVal, ranges[len(ranges)-1][1]) - for i := range ranges { - require.Less(t, ranges[i][0], ranges[i][1]) - if i+1 < len(ranges) { - require.Equal(t, ranges[i][1], ranges[i+1][0]) - } - } -} - -func TestComputeNumericRanges(t *testing.T) { - t.Run("even division", func(t *testing.T) { - ranges := computeNumericRanges(1, 100, 10) - require.Len(t, ranges, 10) - assertContiguousNumeric(t, ranges, 1, 100) - }) - - t.Run("span smaller than partition count", func(t *testing.T) { - ranges := computeNumericRanges(0, 5, 10) - require.Len(t, ranges, 5) - assertContiguousNumeric(t, ranges, 0, 5) - }) - - t.Run("two adjacent values", func(t *testing.T) { - ranges := computeNumericRanges(7, 8, 4) - require.Len(t, ranges, 1) - assertContiguousNumeric(t, ranges, 7, 8) - }) - - t.Run("span does not overflow", func(t *testing.T) { - ranges := computeNumericRanges(math.MinInt64+5, math.MaxInt64-5, 10) - require.Len(t, ranges, 10) - assertContiguousNumeric(t, ranges, math.MinInt64+5, math.MaxInt64-5) - }) - - t.Run("full int64 domain does not overflow", func(t *testing.T) { - ranges := computeNumericRanges(math.MinInt64, math.MaxInt64, 7) - require.Len(t, ranges, 7) - assertContiguousNumeric(t, ranges, math.MinInt64, math.MaxInt64) - }) -} - func TestToRangeFilter(t *testing.T) { t.Run("numeric range is half-open", func(t *testing.T) { filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateNumericPartition(10, 20, false).Range) diff --git a/flow/connectors/utils/partition.go b/flow/connectors/utils/partition.go index 41a12f688..28463541f 100644 --- a/flow/connectors/utils/partition.go +++ b/flow/connectors/utils/partition.go @@ -12,7 +12,6 @@ import ( "google.golang.org/protobuf/types/known/timestamppb" "github.com/PeerDB-io/peerdb/flow/generated/protos" - "github.com/PeerDB-io/peerdb/flow/shared" ) const FullTablePartitionID = "full-table-partition-id" @@ -314,25 +313,22 @@ func (p *PartitionHelper) AddPartitionsWithRange(start any, end any, numPartitio switch r := partition.Range.Range.(type) { case *protos.PartitionRange_IntRange: - size := shared.DivCeil(r.IntRange.End-r.IntRange.Start, numPartitions) - for i := range numPartitions { - if err := p.AddPartition(r.IntRange.Start+size*i, min(r.IntRange.Start+size*(i+1), r.IntRange.End)); err != nil { + for _, rng := range ComputeRanges(r.IntRange.Start, r.IntRange.End, numPartitions) { + if err := p.AddPartition(rng[0], rng[1]); err != nil { return err } } case *protos.PartitionRange_UintRange: - size := shared.DivCeil(r.UintRange.End-r.UintRange.Start, uint64(numPartitions)) - for i := range uint64(numPartitions) { - if err := p.AddPartition(r.UintRange.Start+size*i, min(r.UintRange.Start+size*(i+1), r.UintRange.End)); err != nil { + for _, rng := range ComputeRanges(r.UintRange.Start, r.UintRange.End, numPartitions) { + if err := p.AddPartition(rng[0], rng[1]); err != nil { return err } } case *protos.PartitionRange_TimestampRange: tstart := r.TimestampRange.Start.AsTime().UnixMicro() tend := r.TimestampRange.End.AsTime().UnixMicro() - size := shared.DivCeil(tend-tstart, numPartitions) - for i := range numPartitions { - if err := p.AddPartition(time.UnixMicro(tstart+size*i), time.UnixMicro(min(tstart+size*(i+1), tend))); err != nil { + for _, rng := range ComputeRanges(tstart, tend, numPartitions) { + if err := p.AddPartition(time.UnixMicro(rng[0]), time.UnixMicro(rng[1])); err != nil { return err } } @@ -340,6 +336,27 @@ func (p *PartitionHelper) AddPartitionsWithRange(start any, end any, numPartitio return nil } +func ComputeRanges[T int64 | uint64](minVal T, maxVal T, numPartitions int64) [][2]T { + // span arithmetic is done in uint64 to avoid overflow + span := uint64(maxVal) - uint64(minVal) + step := span / uint64(numPartitions) + if span%uint64(numPartitions) != 0 { + step++ + } + + ranges := make([][2]T, 0, numPartitions) + start := minVal + for { + if remaining := uint64(maxVal) - uint64(start); remaining <= step { + ranges = append(ranges, [2]T{start, maxVal}) + return ranges + } + end := T(uint64(start) + step) + ranges = append(ranges, [2]T{start, end}) + start = end + } +} + func (p *PartitionHelper) getPartitionForStartAndEnd(start any, end any) (*protos.QRepPartition, error) { if start == nil || end == nil { return nil, nil diff --git a/flow/connectors/utils/partition_test.go b/flow/connectors/utils/partition_test.go new file mode 100644 index 000000000..cbbe973e3 --- /dev/null +++ b/flow/connectors/utils/partition_test.go @@ -0,0 +1,138 @@ +package utils + +import ( + "log/slog" + "math" + "testing" + + "github.com/stretchr/testify/require" + "go.temporal.io/sdk/log" + + "github.com/PeerDB-io/peerdb/flow/generated/protos" +) + +func intRangeOf(t *testing.T, partition *protos.QRepPartition) *protos.IntPartitionRange { + t.Helper() + r, ok := partition.Range.Range.(*protos.PartitionRange_IntRange) + require.True(t, ok) + return r.IntRange +} + +func uintRangeOf(t *testing.T, partition *protos.QRepPartition) *protos.UIntPartitionRange { + t.Helper() + r, ok := partition.Range.Range.(*protos.PartitionRange_UintRange) + require.True(t, ok) + return r.UintRange +} + +func assertContiguousRanges[T int64 | uint64](t *testing.T, ranges [][2]T, minVal, maxVal T) { + t.Helper() + require.NotEmpty(t, ranges) + require.Equal(t, minVal, ranges[0][0]) + require.Equal(t, maxVal, ranges[len(ranges)-1][1]) + for i := range ranges { + require.Less(t, ranges[i][0], ranges[i][1]) + if i+1 < len(ranges) { + require.Equal(t, ranges[i][1], ranges[i+1][0]) + } + } +} + +func TestComputeRanges(t *testing.T) { + t.Run("even division", func(t *testing.T) { + ranges := ComputeRanges[int64](1, 100, 10) + require.Len(t, ranges, 10) + assertContiguousRanges(t, ranges, int64(1), int64(100)) + }) + + t.Run("span smaller than partition count", func(t *testing.T) { + ranges := ComputeRanges[int64](0, 5, 10) + require.Len(t, ranges, 5) + assertContiguousRanges(t, ranges, int64(0), int64(5)) + }) + + t.Run("two adjacent values", func(t *testing.T) { + ranges := ComputeRanges[int64](7, 8, 4) + require.Len(t, ranges, 1) + assertContiguousRanges(t, ranges, int64(7), int64(8)) + }) + + t.Run("full int64 domain does not overflow", func(t *testing.T) { + ranges := ComputeRanges[int64](math.MinInt64, math.MaxInt64, 7) + require.Len(t, ranges, 7) + assertContiguousRanges(t, ranges, int64(math.MinInt64), int64(math.MaxInt64)) + }) + + t.Run("full uint64 domain does not overflow", func(t *testing.T) { + ranges := ComputeRanges[uint64](0, math.MaxUint64, 7) + require.Len(t, ranges, 7) + assertContiguousRanges(t, ranges, uint64(0), uint64(math.MaxUint64)) + }) + + t.Run("same min/max values", func(t *testing.T) { + ranges := ComputeRanges[int64](7, 7, 4) + require.Len(t, ranges, 1) + require.Equal(t, int64(7), ranges[0][0]) + require.Equal(t, int64(7), ranges[len(ranges)-1][1]) + }) +} + +func TestAddPartitionsWithRange(t *testing.T) { + logger := log.NewStructuredLogger(slog.Default()) + + t.Run("negative to positive int range covers domain contiguously", func(t *testing.T) { + helper := NewPartitionHelper(logger) + require.NoError(t, helper.AddPartitionsWithRange(int64(-100), int64(100), 4)) + + partitions := helper.GetPartitions() + require.Len(t, partitions, 4) + require.Equal(t, int64(-100), intRangeOf(t, partitions[0]).Start) + require.Equal(t, int64(100), intRangeOf(t, partitions[len(partitions)-1]).End) + for i := 1; i < len(partitions); i++ { + require.Equal(t, intRangeOf(t, partitions[i-1]).End+1, intRangeOf(t, partitions[i]).Start) + } + }) + + t.Run("int range spanning full int64 domain does not overflow", func(t *testing.T) { + helper := NewPartitionHelper(logger) + require.NoError(t, helper.AddPartitionsWithRange(int64(math.MinInt64), int64(math.MaxInt64), 7)) + + partitions := helper.GetPartitions() + require.Len(t, partitions, 7) + require.Equal(t, int64(math.MinInt64), intRangeOf(t, partitions[0]).Start) + require.Equal(t, int64(math.MaxInt64), intRangeOf(t, partitions[len(partitions)-1]).End) + for i, partition := range partitions { + r := intRangeOf(t, partition) + require.LessOrEqual(t, r.Start, r.End) + if i > 0 { + require.Equal(t, intRangeOf(t, partitions[i-1]).End+1, r.Start) + } + } + }) + + t.Run("uint range spanning full uint64 domain does not overflow", func(t *testing.T) { + helper := NewPartitionHelper(logger) + require.NoError(t, helper.AddPartitionsWithRange(uint64(0), uint64(math.MaxUint64), 7)) + + partitions := helper.GetPartitions() + require.Len(t, partitions, 7) + require.Equal(t, uint64(0), uintRangeOf(t, partitions[0]).Start) + require.Equal(t, uint64(math.MaxUint64), uintRangeOf(t, partitions[len(partitions)-1]).End) + for i, partition := range partitions { + r := uintRangeOf(t, partition) + require.LessOrEqual(t, r.Start, r.End) + if i > 0 { + require.Equal(t, uintRangeOf(t, partitions[i-1]).End+1, r.Start) + } + } + }) + + t.Run("span smaller than partition count yields fewer partitions", func(t *testing.T) { + helper := NewPartitionHelper(logger) + require.NoError(t, helper.AddPartitionsWithRange(int64(0), int64(5), 10)) + + partitions := helper.GetPartitions() + require.Len(t, partitions, 5) + require.Equal(t, int64(5), intRangeOf(t, partitions[len(partitions)-1]).End) + }) +}