name: Split materialized data
description: |
  Splits materialized dataset into materialized train, eval, and test data splits.

  The materialized dataset generated by the Feature Transform Engine consists of all the splits
  that were combined into the input transform dataset (i.e., train, eval, and test splits).
  This components splits the output materialized dataset into corresponding materialized data splits
  so that the splits can be used by down-stream training or evaluation components.

    Args:
        materialized_data (Dataset):
            Materialized dataset output by the Feature Transform Engine.

    Returns:
        materialized_train_split (MaterializedSplit):
            Path patern to materialized train split.
        materialized_eval_split (MaterializedSplit):
            Path patern to materialized eval split.
        materialized_test_split (MaterializedSplit):
            Path patern to materialized test split.
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:20230322_0325
    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 indices 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
