sdks/typescript/src/apache_beam/runners/direct_runner.ts - beam - Git at Google

 /*
  * 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.
  */

 import * as uuid from "uuid";
 import * as protobufjs from "protobufjs";
 import Long from "long";

 import * as runnerApi from "../proto/beam_runner_api";
 import * as fnApi from "../proto/beam_fn_api";
 import { PTransform, PCollection } from "../proto/beam_runner_api";
 import { ProcessBundleDescriptor } from "../proto/beam_fn_api";
 import { JobState_Enum } from "../proto/beam_job_api";

 import { Pipeline } from "../internal/pipeline";
 import { Root } from "../pvalue";
 import { impulse, groupByKey } from "../transforms/internal";
 import { Runner, PipelineResult } from "./runner";
 import * as worker from "../worker/worker";
 import * as operators from "../worker/operators";
 import { createStateKey } from "../worker/pardo_context";
 import * as state from "../worker/state";
 import { parDo } from "../transforms/pardo";
 import {
   Window,
   GlobalWindow,
   Instant,
   PaneInfo,
   WindowedValue,
 } from "../values";
 import { PaneInfoCoder } from "../coders/standard_coders";
 import { Coder, Context as CoderContext } from "../coders/coders";
 import * as environments from "../internal/environments";
 import { serializeFn, deserializeFn } from "../internal/serialize";

 const SUPPORTED_REQUIREMENTS: string[] = [];

 export function directRunner(options: Object = {}): Runner {
   return new DirectRunner(options);
 }

 export class DirectRunner extends Runner {
   // All the operators for a given pipeline should share the same state.
   // This global mapping allows operators to look up a shared state object for
   // a given pipeline on deserialization.
   static inMemoryStatesRefs: Map<string, InMemoryStateProvider> = new Map();

   constructor(private options: Object = {}) {
     super();
   }

   unsupportedFeatures(pipeline, options: Object = {}): string[] {
     return [...this.unsupportedFeaturesIter(pipeline, options)];
   }

   *unsupportedFeaturesIter(pipeline, options: Object = {}) {
     const proto: runnerApi.Pipeline = pipeline.proto;
     for (const requirement of proto.requirements) {
       if (!SUPPORTED_REQUIREMENTS.includes(requirement)) {
         yield requirement;
       }
     }

     for (const env of Object.values(proto.components!.environments)) {
       if (
         env.urn &&
         env.urn !== environments.TYPESCRIPT_DEFAULT_ENVIRONMENT_URN
       ) {
         yield env.urn;
       }
     }

     for (const windowing of Object.values(
       proto.components!.windowingStrategies
     )) {
       if (
         ![
           runnerApi.MergeStatus_Enum.UNSPECIFIED,
           runnerApi.MergeStatus_Enum.NON_MERGING,
           runnerApi.MergeStatus_Enum.ALREADY_MERGED,
         ].includes(windowing.mergeStatus)
       ) {
         yield "MergeStatus=" + windowing.mergeStatus;
       }
       if (windowing.outputTime !== runnerApi.OutputTime_Enum.END_OF_WINDOW) {
         yield "OutputTime=" + windowing.outputTime;
       }
     }
   }

   async runPipeline(p): Promise<PipelineResult> {
     // console.dir(p.proto, { depth: null });

     const stateProvider = new InMemoryStateProvider();
     const stateCacheRef = uuid.v4();
     DirectRunner.inMemoryStatesRefs.set(stateCacheRef, stateProvider);

     try {
       const proto = rewriteSideInputs(p.proto, stateCacheRef);
       const descriptor: ProcessBundleDescriptor = {
         id: "",
         transforms: proto.components!.transforms,
         pcollections: proto.components!.pcollections,
         windowingStrategies: proto.components!.windowingStrategies,
         coders: proto.components!.coders,
         environments: proto.components!.environments,
       };

       const processor = new worker.BundleProcessor(
         descriptor,
         null!,
         new state.CachingStateProvider(stateProvider),
         [impulse.urn]
       );
       await processor.process("bundle_id");

       return {
         waitUntilFinish: (duration?: number) =>
           Promise.resolve(JobState_Enum.DONE),
       };
     } finally {
       DirectRunner.inMemoryStatesRefs.delete(stateCacheRef);
     }
   }
 }

 // Only to be used in direct runner, as this will fire an element per worker, not per pipeline.
 class DirectImpulseOperator implements operators.IOperator {
   receiver: operators.Receiver;

   constructor(
     transformId: string,
     transform: PTransform,
     context: operators.OperatorContext
   ) {
     this.receiver = context.getReceiver(
       onlyElement(Object.values(transform.outputs))
     );
   }

   process(wvalue: WindowedValue<unknown>) {
     return operators.NonPromise;
   }

   async startBundle() {
     this.receiver.receive({
       value: new Uint8Array(),
       windows: [new GlobalWindow()],
       pane: PaneInfoCoder.ONE_AND_ONLY_FIRING,
       timestamp: Long.fromValue("-9223372036854775"), // TODO: (Cleanup) Pull constant out of proto, or at least as a constant elsewhere.
     });
   }

   async finishBundle() {}
 }

 operators.registerOperator(impulse.urn, DirectImpulseOperator);

 // Only to be used in direct runner, as this will only group within a single bundle.
 // TODO: (Extension) This could be used as a base for the PGBKOperation operator,
 // and then this class could simply invoke that with an unbounded size and the
 // concat-to-list CombineFn.
 class DirectGbkOperator implements operators.IOperator {
   receiver: operators.Receiver;
   groups: Map<string, unknown[]>;
   keyCoder: Coder<unknown>;
   windowCoder: Coder<Window>;

   constructor(
     transformId: string,
     transform: PTransform,
     context: operators.OperatorContext
   ) {
     this.receiver = context.getReceiver(
       onlyElement(Object.values(transform.outputs))
     );
     const inputPc =
       context.descriptor.pcollections[
         onlyElement(Object.values(transform.inputs))
       ];
     this.keyCoder = context.pipelineContext.getCoder(
       context.descriptor.coders[inputPc.coderId].componentCoderIds[0]
     );
     const windowingStrategy =
       context.descriptor.windowingStrategies[inputPc.windowingStrategyId];
     if (
       windowingStrategy.mergeStatus !== runnerApi.MergeStatus_Enum.NON_MERGING
     ) {
       throw new Error("Unsupported non-merging WindowFn: " + windowingStrategy);
     }
     if (
       windowingStrategy.outputTime !== runnerApi.OutputTime_Enum.END_OF_WINDOW
     ) {
       throw new Error(
         "Unsupported windowing output time: " + windowingStrategy
       );
     }
     this.windowCoder = context.pipelineContext.getCoder(
       windowingStrategy.windowCoderId
     );
   }

   process(wvalue: WindowedValue<any>) {
     for (const window of wvalue.windows) {
       const wkey =
         operators.encodeToBase64(window, this.windowCoder) +
         " " +
         operators.encodeToBase64(wvalue.value.key, this.keyCoder);
       if (!this.groups.has(wkey)) {
         this.groups.set(wkey, []);
       }
       this.groups.get(wkey)!.push(wvalue.value.value);
     }
     return operators.NonPromise;
   }

   async startBundle() {
     this.groups = new Map();
   }

   async finishBundle() {
     for (const [wkey, values] of this.groups) {
       const parts = wkey.split(" ");
       const encodedWindow = parts[0];
       const encodedKey = parts[1];
       const window = operators.decodeFromBase64(
         encodedWindow,
         this.windowCoder
       );
       const maybePromise = this.receiver.receive({
         value: {
           key: operators.decodeFromBase64(encodedKey, this.keyCoder),
           value: values,
         },
         windows: [window],
         timestamp: window.maxTimestamp(),
         pane: PaneInfoCoder.ONE_AND_ONLY_FIRING,
       });
       if (maybePromise !== operators.NonPromise) {
         await maybePromise;
       }
     }
     this.groups = null!;
   }
 }

 operators.registerOperator(groupByKey.urn, DirectGbkOperator);

 /**
  * Rewrites the pipeline to be suitable for running as a single "bundle."
  *
  * Given a pipeline of the form
  *
  *    ...          ...
  *     |            |
  * mainPColl    sidePColl
  *     |            |
  *   parDo <-------/
  *     |
  *    ...
  *
  * Rewrites this to be in the form.
  *
  *    ...          ...
  *     |            |
  * mainPColl    sidePColl
  *     |            |
  *     |      collectSideOp
  *     |            |
  *     |        emptyPColl
  *     | <- - - - -/
  *     |
  *  bufferOp
  *     |
  * mainPcollCopy
  *     |
  *   parDo <----- sidePCollCopy
  *
  * Where collectSideOp collects its inputs for future retrieval and bufferOp
  * buffers all its inputs, invoking the downstream parDo only once finishBundle
  * is called. The dotted line from collectSideOp to bufferOp is not used to
  * pass any data, but acts as a control edge forcing bufferOp to be finished
  * only after all side inputs have been fully collected.
  *
  * A sidePCollCopy is also placed in the graph to allow for lookup of metadata
  * when executing the parDo operaton.
  */
 function rewriteSideInputs(p: runnerApi.Pipeline, pipelineStateRef: string) {
   function uniqueName(container, base) {
     let name = base;
     let counter = 0;
     while (container[name]) {
       name = base + counter++;
     }
     return name;
   }

   p = runnerApi.Pipeline.clone(p);
   const pcolls = p.components!.pcollections;
   const transforms = p.components!.transforms;
   for (const [transformId, transform] of Object.entries(transforms)) {
     if (
       transform.spec?.urn === parDo.urn &&
       Object.keys(transform.inputs).length > 1
     ) {
       const spec = runnerApi.ParDoPayload.fromBinary(transform.spec!.payload);

       // Figure out which input is the main input.
       // TODO: (Typescript) Is there a clean way to do a set difference?
       // There is a proposal for Set.prototype.intersect: https://github.com/tc39/proposal-set-methods
       // But I don't think there is a nicer way at the moment.
       let mainPCollTag: string = undefined!;
       for (const tag of Object.keys(transform.inputs)) {
         if (!spec.sideInputs[tag]) {
           if (mainPCollTag) throw new Error("Multiple non-side inputs.");
           mainPCollTag = tag;
         }
       }
       if (!mainPCollTag) throw new Error("Missing non-side input.");

       // Add the extra logic for each of the side inputs.
       const bufferInputs = { mainPCollTag: transform.inputs[mainPCollTag] };
       for (const side of Object.keys(spec.sideInputs)) {
         const sidePCollId = transform.inputs[side];
         const sideCopyId = uniqueName(pcolls, sidePCollId + "-" + side);
         pcolls[sideCopyId] = pcolls[sidePCollId];
         transform.inputs[side] = sideCopyId;
         const controlPCollId = uniqueName(
           pcolls,
           sidePCollId + "-" + side + "-control"
         );
         pcolls[controlPCollId] = pcolls[transform.inputs[mainPCollTag]];
         bufferInputs[side] = controlPCollId;
         const collectTransformId = uniqueName(
           transforms,
           transformId + "-" + side + "-collect"
         );
         transforms[collectTransformId] = runnerApi.PTransform.create({
           spec: {
             urn: CollectSideOperator.urn,
             payload: serializeFn({
               transformId: transformId,
               sideInputId: side,
               accessPattern: spec.sideInputs[side].accessPattern!.urn,
               pipelineStateRef: pipelineStateRef,
             }),
           },
           inputs: { input: sidePCollId },
           outputs: { out: controlPCollId },
         });
       }

       // Add the buffering operator.
       const bufferId = uniqueName(transforms, transformId + "-buffer");
       const bufferOutputId = uniqueName(pcolls, bufferId + "-out");
       pcolls[bufferOutputId] = pcolls[mainPCollTag];
       transforms[bufferId] = runnerApi.PTransform.create({
         spec: { urn: BufferOperator.urn, payload: new Uint8Array() },
         inputs: bufferInputs,
         outputs: { out: bufferOutputId },
       });
       transform.inputs[mainPCollTag] = bufferOutputId;
     }
   }

   return p;
 }

 class CollectSideOperator implements operators.IOperator {
   static urn = "beam:transforms:node_ts_direct:collect_side:v1";

   stateProvider: InMemoryStateProvider;

   parDoTransformId: string;
   accessPattern: string;
   sideInputId: string;

   receiver: operators.Receiver;
   windowCoder: Coder<Window>;
   elementCoder: Coder<unknown>;

   constructor(
     transformId: string,
     transform: PTransform,
     context: operators.OperatorContext
   ) {
     this.receiver = context.getReceiver(
       onlyElement(Object.values(transform.outputs))
     );
     const payload = deserializeFn(transform.spec!.payload!);
     this.parDoTransformId = payload.transformId;
     this.accessPattern = payload.accessPattern;
     this.sideInputId = payload.sideInputId;
     this.stateProvider = DirectRunner.inMemoryStatesRefs.get(
       payload.pipelineStateRef
     )!;

     const inputPc =
       context.descriptor.pcollections[
         onlyElement(Object.values(transform.inputs))
       ];
     this.elementCoder = context.pipelineContext.getCoder(inputPc.coderId);
     const windowingStrategy =
       context.descriptor.windowingStrategies[inputPc.windowingStrategyId];
     this.windowCoder = context.pipelineContext.getCoder(
       windowingStrategy.windowCoderId
     );
   }

   async startBundle() {}

   process(wvalue: WindowedValue<unknown>) {
     for (const window of wvalue.windows) {
       const writer = new protobufjs.Writer();
       this.elementCoder.encode(
         wvalue.value,
         writer,
         CoderContext.needsDelimiters
       );
       const encodedElement = writer.finish();
       this.stateProvider.appendState(
         createStateKey(
           this.parDoTransformId,
           this.accessPattern,
           this.sideInputId,
           window,
           this.windowCoder
         ),
         encodedElement
       );
     }
     return operators.NonPromise;
   }

   async finishBundle() {}
 }

 operators.registerOperator(CollectSideOperator.urn, CollectSideOperator);

 class BufferOperator implements operators.IOperator {
   static urn = "beam:transforms:node_ts_direct:buffer:v1";
   receiver: operators.Receiver;
   elements: WindowedValue<unknown>[];

   constructor(
     transformId: string,
     transform: PTransform,
     context: operators.OperatorContext
   ) {
     this.receiver = context.getReceiver(
       onlyElement(Object.values(transform.outputs))
     );
   }

   process(wvalue: WindowedValue<unknown>) {
     this.elements.push(wvalue);
     return operators.NonPromise;
   }

   async startBundle() {
     this.elements = [];
   }

   async finishBundle() {
     for (const element of this.elements) {
       const maybePromise = this.receiver.receive(element);
       if (maybePromise !== operators.NonPromise) {
         await maybePromise;
       }
     }
   }
 }

 operators.registerOperator(BufferOperator.urn, BufferOperator);

 class InMemoryStateProvider implements state.StateProvider {
   chunks: Map<string, Uint8Array[]> = new Map();

   getState<T>(
     stateKey: fnApi.StateKey,
     decode: (data: Uint8Array) => T
   ): state.MaybePromise<T> {
     return {
       type: "value",
       value: decode(state.Uint8ArrayConcat(this.getStateEntry(stateKey))),
     };
   }

   appendState(stateKey: fnApi.StateKey, encodedData: Uint8Array) {
     this.getStateEntry(stateKey).push(encodedData);
   }

   getStateEntry<T>(stateKey: fnApi.StateKey) {
     const cacheKey = Buffer.from(fnApi.StateKey.toBinary(stateKey)).toString(
       "base64"
     );
     if (!this.chunks.has(cacheKey)) {
       this.chunks.set(cacheKey, []);
     }
     return this.chunks.get(cacheKey)!;
   }
 }

 /////

 function onlyElement<T>(arg: T[]): T {
   if (arg.length > 1) {
     Error("Expecting exactly one element.");
   }
   return arg[0];
 }
	/*
	* 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.
	*/

	import * as uuid from "uuid";
	import * as protobufjs from "protobufjs";
	import Long from "long";

	import * as runnerApi from "../proto/beam_runner_api";
	import * as fnApi from "../proto/beam_fn_api";
	import { PTransform, PCollection } from "../proto/beam_runner_api";
	import { ProcessBundleDescriptor } from "../proto/beam_fn_api";
	import { JobState_Enum } from "../proto/beam_job_api";

	import { Pipeline } from "../internal/pipeline";
	import { Root } from "../pvalue";
	import { impulse, groupByKey } from "../transforms/internal";
	import { Runner, PipelineResult } from "./runner";
	import * as worker from "../worker/worker";
	import * as operators from "../worker/operators";
	import { createStateKey } from "../worker/pardo_context";
	import * as state from "../worker/state";
	import { parDo } from "../transforms/pardo";
	import {
	Window,
	GlobalWindow,
	Instant,
	PaneInfo,
	WindowedValue,
	} from "../values";
	import { PaneInfoCoder } from "../coders/standard_coders";
	import { Coder, Context as CoderContext } from "../coders/coders";
	import * as environments from "../internal/environments";
	import { serializeFn, deserializeFn } from "../internal/serialize";

	const SUPPORTED_REQUIREMENTS: string[] = [];

	export function directRunner(options: Object = {}): Runner {
	return new DirectRunner(options);
	}

	export class DirectRunner extends Runner {
	// All the operators for a given pipeline should share the same state.
	// This global mapping allows operators to look up a shared state object for
	// a given pipeline on deserialization.
	static inMemoryStatesRefs: Map<string, InMemoryStateProvider> = new Map();

	constructor(private options: Object = {}) {
	super();
	}

	unsupportedFeatures(pipeline, options: Object = {}): string[] {
	return [...this.unsupportedFeaturesIter(pipeline, options)];
	}

	*unsupportedFeaturesIter(pipeline, options: Object = {}) {
	const proto: runnerApi.Pipeline = pipeline.proto;
	for (const requirement of proto.requirements) {
	if (!SUPPORTED_REQUIREMENTS.includes(requirement)) {
	yield requirement;
	}
	}

	for (const env of Object.values(proto.components!.environments)) {
	if (
	env.urn &&
	env.urn !== environments.TYPESCRIPT_DEFAULT_ENVIRONMENT_URN
	) {
	yield env.urn;
	}
	}

	for (const windowing of Object.values(
	proto.components!.windowingStrategies
	)) {
	if (
	![
	runnerApi.MergeStatus_Enum.UNSPECIFIED,
	runnerApi.MergeStatus_Enum.NON_MERGING,
	runnerApi.MergeStatus_Enum.ALREADY_MERGED,
	].includes(windowing.mergeStatus)
	) {
	yield "MergeStatus=" + windowing.mergeStatus;
	}
	if (windowing.outputTime !== runnerApi.OutputTime_Enum.END_OF_WINDOW) {
	yield "OutputTime=" + windowing.outputTime;
	}
	}
	}

	async runPipeline(p): Promise<PipelineResult> {
	// console.dir(p.proto, { depth: null });

	const stateProvider = new InMemoryStateProvider();
	const stateCacheRef = uuid.v4();
	DirectRunner.inMemoryStatesRefs.set(stateCacheRef, stateProvider);

	try {
	const proto = rewriteSideInputs(p.proto, stateCacheRef);
	const descriptor: ProcessBundleDescriptor = {
	id: "",
	transforms: proto.components!.transforms,
	pcollections: proto.components!.pcollections,
	windowingStrategies: proto.components!.windowingStrategies,
	coders: proto.components!.coders,
	environments: proto.components!.environments,
	};

	const processor = new worker.BundleProcessor(
	descriptor,
	null!,
	new state.CachingStateProvider(stateProvider),
	[impulse.urn]
	);
	await processor.process("bundle_id");

	return {
	waitUntilFinish: (duration?: number) =>
	Promise.resolve(JobState_Enum.DONE),
	};
	} finally {
	DirectRunner.inMemoryStatesRefs.delete(stateCacheRef);
	}
	}
	}

	// Only to be used in direct runner, as this will fire an element per worker, not per pipeline.
	class DirectImpulseOperator implements operators.IOperator {
	receiver: operators.Receiver;

	constructor(
	transformId: string,
	transform: PTransform,
	context: operators.OperatorContext
	) {
	this.receiver = context.getReceiver(
	onlyElement(Object.values(transform.outputs))
	);
	}

	process(wvalue: WindowedValue<unknown>) {
	return operators.NonPromise;
	}

	async startBundle() {
	this.receiver.receive({
	value: new Uint8Array(),
	windows: [new GlobalWindow()],
	pane: PaneInfoCoder.ONE_AND_ONLY_FIRING,
	timestamp: Long.fromValue("-9223372036854775"), // TODO: (Cleanup) Pull constant out of proto, or at least as a constant elsewhere.
	});
	}

	async finishBundle() {}
	}

	operators.registerOperator(impulse.urn, DirectImpulseOperator);

	// Only to be used in direct runner, as this will only group within a single bundle.
	// TODO: (Extension) This could be used as a base for the PGBKOperation operator,
	// and then this class could simply invoke that with an unbounded size and the
	// concat-to-list CombineFn.
	class DirectGbkOperator implements operators.IOperator {
	receiver: operators.Receiver;
	groups: Map<string, unknown[]>;
	keyCoder: Coder<unknown>;
	windowCoder: Coder<Window>;

	constructor(
	transformId: string,
	transform: PTransform,
	context: operators.OperatorContext
	) {
	this.receiver = context.getReceiver(
	onlyElement(Object.values(transform.outputs))
	);
	const inputPc =
	context.descriptor.pcollections[
	onlyElement(Object.values(transform.inputs))
	];
	this.keyCoder = context.pipelineContext.getCoder(
	context.descriptor.coders[inputPc.coderId].componentCoderIds[0]
	);
	const windowingStrategy =
	context.descriptor.windowingStrategies[inputPc.windowingStrategyId];
	if (
	windowingStrategy.mergeStatus !== runnerApi.MergeStatus_Enum.NON_MERGING
	) {
	throw new Error("Unsupported non-merging WindowFn: " + windowingStrategy);
	}
	if (
	windowingStrategy.outputTime !== runnerApi.OutputTime_Enum.END_OF_WINDOW
	) {
	throw new Error(
	"Unsupported windowing output time: " + windowingStrategy
	);
	}
	this.windowCoder = context.pipelineContext.getCoder(
	windowingStrategy.windowCoderId
	);
	}

	process(wvalue: WindowedValue<any>) {
	for (const window of wvalue.windows) {
	const wkey =
	operators.encodeToBase64(window, this.windowCoder) +
	" " +
	operators.encodeToBase64(wvalue.value.key, this.keyCoder);
	if (!this.groups.has(wkey)) {
	this.groups.set(wkey, []);
	}
	this.groups.get(wkey)!.push(wvalue.value.value);
	}
	return operators.NonPromise;
	}

	async startBundle() {
	this.groups = new Map();
	}

	async finishBundle() {
	for (const [wkey, values] of this.groups) {
	const parts = wkey.split(" ");
	const encodedWindow = parts[0];
	const encodedKey = parts[1];
	const window = operators.decodeFromBase64(
	encodedWindow,
	this.windowCoder
	);
	const maybePromise = this.receiver.receive({
	value: {
	key: operators.decodeFromBase64(encodedKey, this.keyCoder),
	value: values,
	},
	windows: [window],
	timestamp: window.maxTimestamp(),
	pane: PaneInfoCoder.ONE_AND_ONLY_FIRING,
	});
	if (maybePromise !== operators.NonPromise) {
	await maybePromise;
	}
	}
	this.groups = null!;
	}
	}

	operators.registerOperator(groupByKey.urn, DirectGbkOperator);

	/**
	* Rewrites the pipeline to be suitable for running as a single "bundle."
	*
	* Given a pipeline of the form
	*
	* ... ...
	* \| \|
	* mainPColl sidePColl
	* \| \|
	* parDo <-------/
	* \|
	* ...
	*
	* Rewrites this to be in the form.
	*
	* ... ...
	* \| \|
	* mainPColl sidePColl
	* \| \|
	* \| collectSideOp
	* \| \|
	* \| emptyPColl
	* \| <- - - - -/
	* \|
	* bufferOp
	* \|
	* mainPcollCopy
	* \|
	* parDo <----- sidePCollCopy
	*
	* Where collectSideOp collects its inputs for future retrieval and bufferOp
	* buffers all its inputs, invoking the downstream parDo only once finishBundle
	* is called. The dotted line from collectSideOp to bufferOp is not used to
	* pass any data, but acts as a control edge forcing bufferOp to be finished
	* only after all side inputs have been fully collected.
	*
	* A sidePCollCopy is also placed in the graph to allow for lookup of metadata
	* when executing the parDo operaton.
	*/
	function rewriteSideInputs(p: runnerApi.Pipeline, pipelineStateRef: string) {
	function uniqueName(container, base) {
	let name = base;
	let counter = 0;
	while (container[name]) {
	name = base + counter++;
	}
	return name;
	}

	p = runnerApi.Pipeline.clone(p);
	const pcolls = p.components!.pcollections;
	const transforms = p.components!.transforms;
	for (const [transformId, transform] of Object.entries(transforms)) {
	if (
	transform.spec?.urn === parDo.urn &&
	Object.keys(transform.inputs).length > 1
	) {
	const spec = runnerApi.ParDoPayload.fromBinary(transform.spec!.payload);

	// Figure out which input is the main input.
	// TODO: (Typescript) Is there a clean way to do a set difference?
	// There is a proposal for Set.prototype.intersect: https://github.com/tc39/proposal-set-methods
	// But I don't think there is a nicer way at the moment.
	let mainPCollTag: string = undefined!;
	for (const tag of Object.keys(transform.inputs)) {
	if (!spec.sideInputs[tag]) {
	if (mainPCollTag) throw new Error("Multiple non-side inputs.");
	mainPCollTag = tag;
	}
	}
	if (!mainPCollTag) throw new Error("Missing non-side input.");

	// Add the extra logic for each of the side inputs.
	const bufferInputs = { mainPCollTag: transform.inputs[mainPCollTag] };
	for (const side of Object.keys(spec.sideInputs)) {
	const sidePCollId = transform.inputs[side];
	const sideCopyId = uniqueName(pcolls, sidePCollId + "-" + side);
	pcolls[sideCopyId] = pcolls[sidePCollId];
	transform.inputs[side] = sideCopyId;
	const controlPCollId = uniqueName(
	pcolls,
	sidePCollId + "-" + side + "-control"
	);
	pcolls[controlPCollId] = pcolls[transform.inputs[mainPCollTag]];
	bufferInputs[side] = controlPCollId;
	const collectTransformId = uniqueName(
	transforms,
	transformId + "-" + side + "-collect"
	);
	transforms[collectTransformId] = runnerApi.PTransform.create({
	spec: {
	urn: CollectSideOperator.urn,
	payload: serializeFn({
	transformId: transformId,
	sideInputId: side,
	accessPattern: spec.sideInputs[side].accessPattern!.urn,
	pipelineStateRef: pipelineStateRef,
	}),
	},
	inputs: { input: sidePCollId },
	outputs: { out: controlPCollId },
	});
	}

	// Add the buffering operator.
	const bufferId = uniqueName(transforms, transformId + "-buffer");
	const bufferOutputId = uniqueName(pcolls, bufferId + "-out");
	pcolls[bufferOutputId] = pcolls[mainPCollTag];
	transforms[bufferId] = runnerApi.PTransform.create({
	spec: { urn: BufferOperator.urn, payload: new Uint8Array() },
	inputs: bufferInputs,
	outputs: { out: bufferOutputId },
	});
	transform.inputs[mainPCollTag] = bufferOutputId;
	}
	}

	return p;
	}

	class CollectSideOperator implements operators.IOperator {
	static urn = "beam:transforms:node_ts_direct:collect_side:v1";

	stateProvider: InMemoryStateProvider;

	parDoTransformId: string;
	accessPattern: string;
	sideInputId: string;

	receiver: operators.Receiver;
	windowCoder: Coder<Window>;
	elementCoder: Coder<unknown>;

	constructor(
	transformId: string,
	transform: PTransform,
	context: operators.OperatorContext
	) {
	this.receiver = context.getReceiver(
	onlyElement(Object.values(transform.outputs))
	);
	const payload = deserializeFn(transform.spec!.payload!);
	this.parDoTransformId = payload.transformId;
	this.accessPattern = payload.accessPattern;
	this.sideInputId = payload.sideInputId;
	this.stateProvider = DirectRunner.inMemoryStatesRefs.get(
	payload.pipelineStateRef
	)!;

	const inputPc =
	context.descriptor.pcollections[
	onlyElement(Object.values(transform.inputs))
	];
	this.elementCoder = context.pipelineContext.getCoder(inputPc.coderId);
	const windowingStrategy =
	context.descriptor.windowingStrategies[inputPc.windowingStrategyId];
	this.windowCoder = context.pipelineContext.getCoder(
	windowingStrategy.windowCoderId
	);
	}

	async startBundle() {}

	process(wvalue: WindowedValue<unknown>) {
	for (const window of wvalue.windows) {
	const writer = new protobufjs.Writer();
	this.elementCoder.encode(
	wvalue.value,
	writer,
	CoderContext.needsDelimiters
	);
	const encodedElement = writer.finish();
	this.stateProvider.appendState(
	createStateKey(
	this.parDoTransformId,
	this.accessPattern,
	this.sideInputId,
	window,
	this.windowCoder
	),
	encodedElement
	);
	}
	return operators.NonPromise;
	}

	async finishBundle() {}
	}

	operators.registerOperator(CollectSideOperator.urn, CollectSideOperator);

	class BufferOperator implements operators.IOperator {
	static urn = "beam:transforms:node_ts_direct:buffer:v1";
	receiver: operators.Receiver;
	elements: WindowedValue<unknown>[];

	constructor(
	transformId: string,
	transform: PTransform,
	context: operators.OperatorContext
	) {
	this.receiver = context.getReceiver(
	onlyElement(Object.values(transform.outputs))
	);
	}

	process(wvalue: WindowedValue<unknown>) {
	this.elements.push(wvalue);
	return operators.NonPromise;
	}

	async startBundle() {
	this.elements = [];
	}

	async finishBundle() {
	for (const element of this.elements) {
	const maybePromise = this.receiver.receive(element);
	if (maybePromise !== operators.NonPromise) {
	await maybePromise;
	}
	}
	}
	}

	operators.registerOperator(BufferOperator.urn, BufferOperator);

	class InMemoryStateProvider implements state.StateProvider {
	chunks: Map<string, Uint8Array[]> = new Map();

	getState<T>(
	stateKey: fnApi.StateKey,
	decode: (data: Uint8Array) => T
	): state.MaybePromise<T> {
	return {
	type: "value",
	value: decode(state.Uint8ArrayConcat(this.getStateEntry(stateKey))),
	};
	}

	appendState(stateKey: fnApi.StateKey, encodedData: Uint8Array) {
	this.getStateEntry(stateKey).push(encodedData);
	}

	getStateEntry<T>(stateKey: fnApi.StateKey) {
	const cacheKey = Buffer.from(fnApi.StateKey.toBinary(stateKey)).toString(
	"base64"
	);
	if (!this.chunks.has(cacheKey)) {
	this.chunks.set(cacheKey, []);
	}
	return this.chunks.get(cacheKey)!;
	}
	}

	/////

	function onlyElement<T>(arg: T[]): T {
	if (arg.length > 1) {
	Error("Expecting exactly one element.");
	}
	return arg[0];
	}