44 "cmp"
55 "fmt"
66 "log/slog"
7+ "math/big"
8+ "regexp"
9+ "strings"
710 "time"
811
912 "github.com/google/uuid"
@@ -15,8 +18,21 @@ import (
1518 "github.com/PeerDB-io/peerdb/flow/shared"
1619)
1720
21+ var (
22+ UuidLowerRe = regexp .MustCompile (`^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$` )
23+ UuidUpperRe = regexp .MustCompile (`^[0-9A-F]{8}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{4}-[0-9A-F]{12}$` )
24+ )
25+
1826const FullTablePartitionID = "full-table-partition-id"
1927
28+ func FullTablePartition () []* protos.QRepPartition {
29+ return []* protos.QRepPartition {{
30+ PartitionId : FullTablePartitionID ,
31+ Range : nil ,
32+ FullTablePartition : true ,
33+ }}
34+ }
35+
2036type PartitionRangeType string
2137
2238const (
@@ -141,6 +157,21 @@ func createUIntPartition(start uint64, end uint64) *protos.QRepPartition {
141157 }
142158}
143159
160+ func createStringPartition (start string , end string , endInclusive bool ) * protos.QRepPartition {
161+ return & protos.QRepPartition {
162+ PartitionId : uuid .NewString (),
163+ Range : & protos.PartitionRange {
164+ Range : & protos.PartitionRange_StringRange {
165+ StringRange : & protos.StringPartitionRange {
166+ Start : start ,
167+ End : end ,
168+ EndInclusive : endInclusive ,
169+ },
170+ },
171+ },
172+ }
173+ }
174+
144175type PartitionHelper struct {
145176 logger log.Logger
146177 prevStart any
@@ -328,10 +359,97 @@ func (p *PartitionHelper) AddPartitionsWithRange(start any, end any, numPartitio
328359 return err
329360 }
330361 }
362+ case * protos.PartitionRange_StringRange :
363+ return p .addStringPartitions (r .StringRange .Start , r .StringRange .End , numPartitions )
331364 }
332365 return nil
333366}
334367
368+ type casing int
369+
370+ const (
371+ unknown casing = iota
372+ lower
373+ upper
374+ )
375+
376+ func (p * PartitionHelper ) addStringPartitions (minVal string , maxVal string , numPartitions int64 ) error {
377+ isUuid , hexCasing := detectUuidWithHexCasing (minVal , maxVal )
378+
379+ if ! isUuid {
380+ // TODO: implement parallel snapshotting for arbitrary strings
381+ p .logger .Info ("string watermark column is not uuid, falling back to full table partition" )
382+ p .partitions = append (p .partitions , FullTablePartition ()... )
383+ return nil
384+ }
385+
386+ // The regex matching guarantees UUIDs, so error should not happen
387+ minInt , err := uuidToBigInt (minVal )
388+ if err != nil {
389+ return fmt .Errorf ("failed to parse min watermark as UUID: %w" , err )
390+ }
391+ maxInt , err := uuidToBigInt (maxVal )
392+ if err != nil {
393+ return fmt .Errorf ("failed to parse max watermark as UUID: %w" , err )
394+ }
395+
396+ start := minVal
397+ step := shared .BigIntDivCeil (new (big.Int ).Sub (maxInt , minInt ), big .NewInt (numPartitions ))
398+ for value := new (big.Int ).Add (minInt , step ); value .Cmp (maxInt ) < 0 ; value .Add (value , step ) {
399+ end , err := bigIntToUUID (value , hexCasing )
400+ if err != nil {
401+ return fmt .Errorf ("failed to render partition boundary: %w" , err )
402+ }
403+ p .partitions = append (p .partitions , createStringPartition (start , end , false ))
404+ start = end
405+ }
406+ p .partitions = append (p .partitions , createStringPartition (start , maxVal , true ))
407+ return nil
408+ }
409+
410+ // detectUuidWithHexCasing best-effort classifies a string watermark column by
411+ // inspecting only its min and max bounds: it determines whether both are
412+ // canonical UUIDs, and if so, the shared casing of their hex letters.
413+ //
414+ // Because only the bounds are examined, the classification can be wrong in
415+ // the following rare cases:
416+ // - non-UUID rows can exist between UUID-shaped min/max bounds;
417+ // - non-boundary UUIDs contain a mix of upper and lower case
418+ // - min/max bounds both don't contain hex, so casing blindly defaults to lower
419+
420+ // These cases may lead to partition skew, but will not affect correctness.
421+ func detectUuidWithHexCasing (minVal string , maxVal string ) (bool , casing ) {
422+ switch {
423+ case UuidLowerRe .MatchString (minVal ) && UuidLowerRe .MatchString (maxVal ):
424+ return true , lower
425+ case UuidUpperRe .MatchString (minVal ) && UuidUpperRe .MatchString (maxVal ):
426+ return true , upper
427+ default :
428+ return false , unknown
429+ }
430+ }
431+
432+ func uuidToBigInt (s string ) (* big.Int , error ) {
433+ u , err := uuid .Parse (s )
434+ if err != nil {
435+ return nil , err
436+ }
437+ return new (big.Int ).SetBytes (u [:]), nil
438+ }
439+
440+ func bigIntToUUID (n * big.Int , casing casing ) (string , error ) {
441+ if n .BitLen () > 128 {
442+ return "" , fmt .Errorf ("value does not fit in a UUID (%d bits)" , n .BitLen ())
443+ }
444+ var b [16 ]byte
445+ n .FillBytes (b [:])
446+ s := uuid .UUID (b ).String ()
447+ if casing == upper {
448+ s = strings .ToUpper (s )
449+ }
450+ return s , nil
451+ }
452+
335453func (p * PartitionHelper ) getPartitionForStartAndEnd (start any , end any ) (* protos.QRepPartition , error ) {
336454 if start == nil || end == nil {
337455 return nil , nil
@@ -357,6 +475,8 @@ func (p *PartitionHelper) getPartitionForStartAndEnd(start any, end any) (*proto
357475 return createTimePartition (v , end .(time.Time )), nil
358476 case pgtype.TID :
359477 return createTIDPartition (v , end .(pgtype.TID )), nil
478+ case string :
479+ return createStringPartition (v , end .(string ), true ), nil
360480 default :
361481 return nil , fmt .Errorf ("unsupported type: %T" , v )
362482 }
0 commit comments