[apache#24789][prism] Handlers for combine, ParDo, GBK, Flatten (apac…

…he#25558)

* [prism] add in handlers

* [prism] executor interface

* delint - rm processor

* [prism] comments

sdks/go/pkg/beam/runners/prism/internal/execute.go

@@ -15,8 +15,20 @@
 
 package internal
 
+import (
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/graph/mtime"
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/worker"
+)
+
 // stage represents a fused subgraph.
 // temporary implementation to break up PRs.
 type stage struct {
 	transforms []string
 }
+
+type transformExecuter interface {
+	ExecuteUrns() []string
+	ExecuteWith(t *pipepb.PTransform) string
+	ExecuteTransform(tid string, t *pipepb.PTransform, comps *pipepb.Components, watermark mtime.Time, data [][]byte) *worker.B
+}

sdks/go/pkg/beam/runners/prism/internal/handlecombine.go

@@ -0,0 +1,209 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements.  See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package internal
+
+import (
+	"fmt"
+	"reflect"
+
+	pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1"
+	"github.com/apache/beam/sdks/v2/go/pkg/beam/runners/prism/internal/urns"
+	"google.golang.org/protobuf/proto"
+)
+
+// This file retains the logic for the combine handler
+
+// CombineCharacteristic holds the configuration for Combines.
+type CombineCharacteristic struct {
+	EnableLifting bool // Sets whether a combine composite does combiner lifting or not.
+}
+
+// TODO figure out the factory we'd like.
+
+func Combine(config any) *combine {
+	return &combine{config: config.(CombineCharacteristic)}
+}
+
+// combine represents an instance of the combine handler.
+type combine struct {
+	config CombineCharacteristic
+}
+
+// ConfigURN returns the name for combine in the configuration file.
+func (*combine) ConfigURN() string {
+	return "combine"
+}
+
+func (*combine) ConfigCharacteristic() reflect.Type {
+	return reflect.TypeOf((*CombineCharacteristic)(nil)).Elem()
+}
+
+var _ transformPreparer = (*combine)(nil)
+
+func (*combine) PrepareUrns() []string {
+	return []string{urns.TransformCombinePerKey}
+}
+
+// PrepareTransform returns lifted combines and removes the leaves if enabled. Otherwise returns nothing.
+func (h *combine) PrepareTransform(tid string, t *pipepb.PTransform, comps *pipepb.Components) (*pipepb.Components, []string) {
+	// If we aren't lifting, the "default impl" for combines should be sufficient.
+	if !h.config.EnableLifting {
+		return nil, nil
+	}
+
+	// To lift a combine, the spec should contain a CombinePayload.
+	// That contains the actual FunctionSpec for the DoFn, and the
+	// id for the accumulator coder.
+	// We can synthetically produce/determine the remaining coders for
+	// the Input and Output types from the existing PCollections.
+	//
+	// This means we also need to synthesize pcollections with the accumulator coder too.
+
+	// What we have:
+	//  Input PCol: KV<K, I>      -- INPUT
+	//  -> GBK := KV<K, Iter<I>>  -- GROUPED_I
+	//  -> Combine := KV<K, O>    -- OUTPUT
+	//
+	// What we want:
+	//  Input PCol: KV<K, I>             -- INPUT
+	//  -> PreCombine := KV<K, A>        -- LIFTED
+	//  -> GBK -> KV<K, Iter<A>>         -- GROUPED_A
+	//  -> MergeAccumulators := KV<K, A> -- MERGED_A
+	//  -> ExtractOutput -> KV<K, O>     -- OUTPUT
+	//
+	// First we need to produce new coders for Iter<A>, KV<K, Iter<A>>, and KV<K, A>.
+	// The A coder ID is in the combine payload.
+	//
+	// Then we can produce the PCollections.
+	// We can reuse the INPUT and OUTPUT PCollections.
+	// We need LIFTED to have KV<K, A>  kv_k_a
+	// We need GROUPED_A to have KV<K, Iter<A>> kv_k_iter_a
+	// We need MERGED_A to have KV<K, A> kv_k_a
+	//
+	// GROUPED_I ends up unused.
+	//
+	// The PCollections inherit the properties of the Input PCollection
+	// such as Boundedness, and Windowing Strategy.
+	//
+	// With these, we can produce the PTransforms with the appropriate URNs for the
+	// different parts of the composite, and return the new components.
+
+	cmbPayload := t.GetSpec().GetPayload()
+	cmb := &pipepb.CombinePayload{}
+	if err := (proto.UnmarshalOptions{}).Unmarshal(cmbPayload, cmb); err != nil {
+		panic(fmt.Sprintf("unable to decode ParDoPayload for transform[%v]", t.GetUniqueName()))
+	}
+
+	// First lets get the key coder ID.
+	var pcolInID string
+	// There's only one input.
+	for _, pcol := range t.GetInputs() {
+		pcolInID = pcol
+	}
+	inputPCol := comps.GetPcollections()[pcolInID]
+	kvkiID := inputPCol.GetCoderId()
+	kID := comps.GetCoders()[kvkiID].GetComponentCoderIds()[0]
+
+	// Now we can start synthesis!
+	// Coder IDs
+	aID := cmb.AccumulatorCoderId
+
+	ckvprefix := "c" + tid + "_kv_"
+
+	iterACID := "c" + tid + "_iter_" + aID
+	kvkaCID := ckvprefix + kID + "_" + aID
+	kvkIterACID := ckvprefix + kID + "_iter" + aID
+
+	// PCollection IDs
+	nprefix := "n" + tid + "_"
+	liftedNID := nprefix + "lifted"
+	groupedNID := nprefix + "grouped"
+	mergedNID := nprefix + "merged"
+
+	// Now we need the output collection ID
+	var pcolOutID string
+	// There's only one input.
+	for _, pcol := range t.GetOutputs() {
+		pcolOutID = pcol
+	}
+
+	// Transform IDs
+	eprefix := "e" + tid + "_"
+	liftEID := eprefix + "lift"
+	gbkEID := eprefix + "gbk"
+	mergeEID := eprefix + "merge"
+	extractEID := eprefix + "extract"
+
+	coder := func(urn string, componentIDs ...string) *pipepb.Coder {
+		return &pipepb.Coder{
+			Spec: &pipepb.FunctionSpec{
+				Urn: urn,
+			},
+			ComponentCoderIds: componentIDs,
+		}
+	}
+
+	pcol := func(name, coderID string) *pipepb.PCollection {
+		return &pipepb.PCollection{
+			UniqueName:          name,
+			CoderId:             coderID,
+			IsBounded:           inputPCol.GetIsBounded(),
+			WindowingStrategyId: inputPCol.GetWindowingStrategyId(),
+		}
+	}
+
+	tform := func(name, urn, in, out, env string) *pipepb.PTransform {
+		return &pipepb.PTransform{
+			UniqueName: name,
+			Spec: &pipepb.FunctionSpec{
+				Urn:     urn,
+				Payload: cmbPayload,
+			},
+			Inputs: map[string]string{
+				"i0": in,
+			},
+			Outputs: map[string]string{
+				"i0": out,
+			},
+			EnvironmentId: env,
+		}
+	}
+
+	newComps := &pipepb.Components{
+		Coders: map[string]*pipepb.Coder{
+			iterACID:    coder(urns.CoderIterable, aID),
+			kvkaCID:     coder(urns.CoderKV, kID, aID),
+			kvkIterACID: coder(urns.CoderKV, kID, iterACID),
+		},
+		Pcollections: map[string]*pipepb.PCollection{
+			liftedNID:  pcol(liftedNID, kvkaCID),
+			groupedNID: pcol(groupedNID, kvkIterACID),
+			mergedNID:  pcol(mergedNID, kvkaCID),
+		},
+		Transforms: map[string]*pipepb.PTransform{
+			liftEID:    tform(liftEID, urns.TransformPreCombine, pcolInID, liftedNID, t.GetEnvironmentId()),
+			gbkEID:     tform(gbkEID, urns.TransformGBK, liftedNID, groupedNID, ""),
+			mergeEID:   tform(mergeEID, urns.TransformMerge, groupedNID, mergedNID, t.GetEnvironmentId()),
+			extractEID: tform(mergeEID, urns.TransformExtract, mergedNID, pcolOutID, t.GetEnvironmentId()),
+		},
+	}
+
+	// Now we return everything!
+	// TODO recurse through sub transforms to remove?
+	// We don't need to remove the composite, since we don't add it in
+	// when we return the new transforms, so it's not in the topology.
+	return newComps, t.GetSubtransforms()
+}