Merge pull request #1 from KeranYang/mapT

Map t

pkg/apis/proto/function/v1/udfunction.proto

@@ -8,10 +8,14 @@ import "google/protobuf/empty.proto";
 package function.v1;
 
 service UserDefinedFunction {
-  // Applies a function to each datum element.
+  // MapFn applies a function to each datum element without modifying event time.
   rpc MapFn(Datum) returns (DatumList);
 
-  // Applies a reduce function to a datum stream.
+  // MapTFn applies a function to each datum element.
+  // In addition to map function, MapTFn also supports assigning a new event time to datum.
+  rpc MapTFn(Datum) returns (DatumList);
+
+  // ReduceFn applies a reduce function to a datum stream.
   rpc ReduceFn(stream Datum) returns (DatumList);
 
   // IsReady is the heartbeat endpoint for gRPC.

pkg/function/client/client.go

@@ -53,7 +53,7 @@ func (c *client) IsReady(ctx context.Context, in *emptypb.Empty) (bool, error) {
 	return resp.GetReady(), nil
 }
 
-// MapFn applies a function to each datum element.
+// MapFn applies a function to each datum element without modifying event time.
 func (c *client) MapFn(ctx context.Context, datum *functionpb.Datum) ([]*functionpb.Datum, error) {
 	mappedDatumList, err := c.grpcClt.MapFn(ctx, datum)
 	if err != nil {
@@ -63,6 +63,17 @@ func (c *client) MapFn(ctx context.Context, datum *functionpb.Datum) ([]*functio
 	return mappedDatumList.GetElements(), nil
 }
 
+// MapTFn applies a function to each datum element.
+// In addition to map function, MapTFn also supports assigning a new event time to datum.
+func (c *client) MapTFn(ctx context.Context, datum *functionpb.Datum) ([]*functionpb.Datum, error) {
+	mappedDatumList, err := c.grpcClt.MapTFn(ctx, datum)
+	if err != nil {
+		return nil, fmt.Errorf("failed to execute c.grpcClt.MapTFn(): %w", err)
+	}
+
+	return mappedDatumList.GetElements(), nil
+}
+
 // ReduceFn applies a reduce function to a datum stream.
 func (c *client) ReduceFn(ctx context.Context, datumStreamCh <-chan *functionpb.Datum) ([]*functionpb.Datum, error) {
 	stream, err := c.grpcClt.ReduceFn(ctx)

pkg/function/clienttest/client_test.go

@@ -93,3 +93,40 @@ func TestMapFn(t *testing.T) {
 	assert.Nil(t, got)
 	assert.EqualError(t, err, "failed to execute c.grpcClt.MapFn(): mock MapFn error")
 }
+
+func TestMapTFn(t *testing.T) {
+	var ctx = context.Background()
+
+	ctrl := gomock.NewController(t)
+	defer ctrl.Finish()
+
+	mockClient := funcmock.NewMockUserDefinedFunctionClient(ctrl)
+	testDatum := &functionpb.Datum{
+		Key:       "test_success_key",
+		Value:     []byte(`forward_message`),
+		EventTime: &functionpb.EventTime{EventTime: timestamppb.New(time.Unix(1661169600, 0))},
+		Watermark: &functionpb.Watermark{Watermark: timestamppb.New(time.Time{})},
+	}
+	mockClient.EXPECT().MapTFn(gomock.Any(), &rpcMsg{msg: testDatum}).Return(&functionpb.DatumList{
+		Elements: []*functionpb.Datum{
+			testDatum,
+		},
+	}, nil)
+	mockClient.EXPECT().MapTFn(gomock.Any(), &rpcMsg{msg: testDatum}).Return(&functionpb.DatumList{
+		Elements: []*functionpb.Datum{
+			nil,
+		},
+	}, fmt.Errorf("mock MapTFn error"))
+	testClient, err := New(mockClient)
+	assert.NoError(t, err)
+	reflect.DeepEqual(testClient, &client{
+		grpcClt: mockClient,
+	})
+	got, err := testClient.MapTFn(ctx, testDatum)
+	reflect.DeepEqual(got, testDatum)
+	assert.NoError(t, err)
+
+	got, err = testClient.MapTFn(ctx, testDatum)
+	assert.Nil(t, got)
+	assert.EqualError(t, err, "failed to execute c.grpcClt.MapTFn(): mock MapTFn error")
+}

pkg/function/clienttest/clienttest.go

@@ -33,7 +33,7 @@ func (c *client) IsReady(ctx context.Context, in *emptypb.Empty) (bool, error) {
 	return resp.GetReady(), nil
 }
 
-// MapFn applies a function to each datum element.
+// MapFn applies a function to each datum element without modifying event time.
 func (c *client) MapFn(ctx context.Context, datum *functionpb.Datum) ([]*functionpb.Datum, error) {
 	mappedDatumList, err := c.grpcClt.MapFn(ctx, datum)
 	if err != nil {
@@ -43,6 +43,17 @@ func (c *client) MapFn(ctx context.Context, datum *functionpb.Datum) ([]*functio
 	return mappedDatumList.GetElements(), nil
 }
 
+// MapTFn applies a function to each datum element.
+// In addition to map function, MapTFn also supports assigning a new event time to datum.
+func (c *client) MapTFn(ctx context.Context, datum *functionpb.Datum) ([]*functionpb.Datum, error) {
+	mappedDatumList, err := c.grpcClt.MapTFn(ctx, datum)
+	if err != nil {
+		return nil, fmt.Errorf("failed to execute c.grpcClt.MapTFn(): %w", err)
+	}
+
+	return mappedDatumList.GetElements(), nil
+}
+
 // ReduceFn applies a reduce function to a datum stream.
 func (c *client) ReduceFn(ctx context.Context, datumStreamCh <-chan *functionpb.Datum) ([]*functionpb.Datum, error) {
 	stream, err := c.grpcClt.ReduceFn(ctx)