name: Split materialized data
description: Splits materialized_data into materialized_data test, train, and eval
  splits.
inputs:
- {name: materialized_data, type: Dataset, description: materialized_data dataset
    output by FTE.}
outputs:
- {name: materialized_train_split, type: MaterializedSplit, description: Path patern
    to materialized_train_split.}
- {name: materialized_eval_split, type: MaterializedSplit, description: Path patern
    to materialized_eval_split.}
- {name: materialized_test_split, type: MaterializedSplit, description: Path patern
    to materialized_test_split.}
implementation:
  container:
    image: us-docker.pkg.dev/vertex-ai/automl-tabular/dataflow-worker:20220629_0725_RC00
    command:
    - sh
    - -ec
    - |
      program_path=$(mktemp -d)
      printf "%s" "$0" > "$program_path/ephemeral_component.py"
      python3 -m kfp.v2.components.executor_main                         --component_module_path                         "$program_path/ephemeral_component.py"                         "$@"
    - |2+

      import kfp
      from kfp.v2 import dsl
      from kfp.v2.dsl import *
      from typing import *

      def _split_materialized_data(
          materialized_data: Input[Dataset],
          materialized_train_split: OutputPath('MaterializedSplit'),
          materialized_eval_split: OutputPath('MaterializedSplit'),
          materialized_test_split: OutputPath('MaterializedSplit')):
        """Splits materialized_data into materialized_data test, train, and eval splits.

        Necessary adapter between FTE pipeline and trainer.

        Args:
          materialized_data: materialized_data dataset output by FTE.
          materialized_train_split: Path patern to materialized_train_split.
          materialized_eval_split: Path patern to materialized_eval_split.
          materialized_test_split: Path patern to materialized_test_split.
        """
        # pylint: disable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported
        import json
        import tensorflow as tf
        # pylint: enable=g-import-not-at-top,import-outside-toplevel,redefined-outer-name,reimported

        with tf.io.gfile.GFile(materialized_data.path, 'r') as f:
          artifact_path = f.read()

        # needed to import tf because this is a path in gs://
        with tf.io.gfile.GFile(artifact_path, 'r') as f:
          materialized_data_json = json.load(f)

        file_patterns = materialized_data_json['tf_record_data_source'][
            'file_patterns']

        # we map indicies to file patterns based on the ordering of insertion order
        # in our transform_data (see above in _generate_analyze_and_transform_data)
        with tf.io.gfile.GFile(materialized_train_split, 'w') as f:
          f.write(file_patterns[0])

        with tf.io.gfile.GFile(materialized_eval_split, 'w') as f:
          f.write(file_patterns[1])

        with tf.io.gfile.GFile(materialized_test_split, 'w') as f:
          f.write(file_patterns[2])

    args:
    - --executor_input
    - {executorInput: null}
    - --function_to_execute
    - _split_materialized_data
