org.flyte.examples.PhoneBookWorkflow Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2021 Flyte Authors
 *
 * Licensed 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 org.flyte.examples;

import com.google.auto.service.AutoService;
import com.google.auto.value.AutoValue;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.flyte.flytekit.SdkBindingData;
import org.flyte.flytekit.SdkBindingDataFactory;
import org.flyte.flytekit.SdkTypes;
import org.flyte.flytekit.SdkWorkflow;
import org.flyte.flytekit.SdkWorkflowBuilder;
import org.flyte.flytekit.jackson.JacksonSdkType;

@AutoService(SdkWorkflow.class)
public class PhoneBookWorkflow extends SdkWorkflow {

  private static final List NAMES = Arrays.asList("frodo", "bilbo");
  private static final Map PHONE_BOOK = new HashMap<>();

  static {
    PHONE_BOOK.put("frodo", "123");
    PHONE_BOOK.put("bilbo", "456");
    PHONE_BOOK.put("gandalf", "789");
  }

  @AutoValue
  public abstract static class Output {
    public abstract SdkBindingData> phoneNumbers();

    /**
     * Wraps the constructor of the generated output value class.
     *
     * @param phoneNumbers the String literal output of {@link NodeMetadataExampleWorkflow}
     * @return output of NodeMetadataExampleWorkflow
     */
    public static PhoneBookWorkflow.Output create(SdkBindingData> phoneNumbers) {
      return new AutoValue_PhoneBookWorkflow_Output(phoneNumbers);
    }
  }

  public PhoneBookWorkflow() {
    super(SdkTypes.nulls(), JacksonSdkType.of(PhoneBookWorkflow.Output.class));
  }

  @Override
  public Output expand(SdkWorkflowBuilder builder, Void noInput) {
    SdkBindingData> phoneBook = SdkBindingDataFactory.ofStringMap(PHONE_BOOK);

    SdkBindingData> searchKeys = SdkBindingDataFactory.ofStringCollection(NAMES);

    SdkBindingData> phoneNumbers =
        builder
            .apply(
                "search",
                new BatchLookUpTask(),
                BatchLookUpTask.Input.create(phoneBook, searchKeys))
            .getOutputs()
            .values();

    return PhoneBookWorkflow.Output.create(phoneNumbers);
  }
}