introducing DefaultPartitions()

Signed-off-by: Yashash H L <[email protected]>

pkg/sourcer/interface.go

@@ -18,7 +18,9 @@ type Sourcer interface {
 	// When the return value is negative, it indicates the pending information is not available.
 	// With pending information being not available, the Numaflow platform doesn't auto-scale the source.
 	Pending(ctx context.Context) int64
-	//Partitions returns the partitions associated with the source.
+	//Partitions returns the partitions associated with the source, will be used by the platform to determine
+	// the partitions to which the watermark should be published. If the source doesn't have partitions,
+	// DefaultPartitions() can be used to return the default partitions.
 	Partitions(ctx context.Context) []int32
 }
 

pkg/sourcer/message.go

@@ -1,6 +1,13 @@
 package sourcer
 
-import "time"
+import (
+	"os"
+	"strconv"
+	"time"
+)
+
+// create default partition id from the environment variable "NUMAFLOW_REPLICA"
+var defaultPartitionId, _ = strconv.Atoi(os.Getenv("NUMAFLOW_REPLICA"))
 
 // Message is used to wrap the data return by UDSource
 type Message struct {
@@ -51,6 +58,19 @@ func NewOffset(value []byte, partitionId int32) Offset {
 	return Offset{value: value, partitionId: partitionId}
 }
 
+// NewOffsetWithDefaultPartitionId creates an Offset with value and default partition id
+func NewOffsetWithDefaultPartitionId(value []byte) Offset {
+	return Offset{value: value, partitionId: DefaultPartitions()[0]}
+}
+
+// DefaultPartitions returns default partitions for the source
+// it can be used in the Partitions() function of the Sourcer implementation
+// if the source doesn't have partitions, default partition will be pod replica
+// index of the source.
+func DefaultPartitions() []int32 {
+	return []int32{int32(defaultPartitionId)}
+}
+
 // Value returns value of the offset
 func (o Offset) Value() []byte {
 	return o.value