docstr: Runs an analysis pipeline with cache.
  
  Args:
    feature_spec: A feature_spec for the input data.
    input_data_dict: Dict[str, List[Dict[str, primitive]]] the input data used
      for analysis.
    preprocessing_fn: The preprocessing_fn used for analysis.
    cache_dict: Dict[str, Dict[str, List[bytes]]], input cache dict. If
      provided, should_read_cache must be False.
    should_read_cache: A bool indicating if the pipeline should read cache. If
      True, cache_dict must be False.
    datasets_to_transform: List[str], list of dataset keys to transform.
    expected_transform_data: List[Dict[str, primitive]], the expected
      transformed data, should be the same for each dataset.
    expected_cache: Dict[str, Dict[str, bytes]], expected encoded cache.
    transform_fn_output_dir: A directory where the output transform_fn should
      be written to, if None provided it will not be written.
    use_tf_compat_v1: If True, TFT's public APIs (e.g. AnalyzeDataset) will
      use Tensorflow in compat.v1 mode. Defaults to `True`.
  
  Returns:
    A _RunPipelineResult.
docstr_tok: ["Runs","an","analysis","pipeline","with","cache",".","Args",":","feature_spec",":","A","feature_spec","for","the","input","data",".","input_data_dict",":","Dict","[","str",",","List","[","Dict","[","str",",","primitive","]","]","]","the","input","data","used","for","analysis",".","preprocessing_fn",":","The","preprocessing_fn","used","for","analysis",".","cache_dict",":","Dict","[","str",",","Dict","[","str",",","List","[","bytes","]","]","]",",","input","cache","dict",".","If","provided",",","should_read_cache","must","be","False",".","should_read_cache",":","A","bool","indicating","if","the","pipeline","should","read","cache",".","If","True",",","cache_dict","must","be","False",".","datasets_to_transform",":","List","[","str","]",",","list","of","dataset","keys","to","transform",".","expected_transform_data",":","List","[","Dict","[","str",",","primitive","]","]",",","the","expected","transformed","data",",","should","be","the","same","for","each","dataset",".","expected_cache",":","Dict","[","str",",","Dict","[","str",",","bytes","]","]",",","expected","encoded","cache",".","transform_fn_output_dir",":","A","directory","where","the","output","transform_fn","should","be","written","to",",","if","None","provided","it","will","not","be","written",".","use_tf_compat_v1",":","If","True",",","TFT","'s","public","APIs","(","e.g",".","AnalyzeDataset",")","will","use","Tensorflow","in","compat.v1","mode",".","Defaults","to","`","True","`",".","Returns",":","A","_RunPipelineResult","."]
code: def _run_pipeline(
      self,
      feature_spec,
      input_data_dict,
      preprocessing_fn,
      cache_dict=None,
      should_read_cache=False,
      datasets_to_transform=None,
      expected_transform_data=None,
      expected_cache=None,
      transform_fn_output_dir="",
      use_tf_compat_v1=True,
  ) -> _RunPipelineResult:
      """Runs an analysis pipeline with cache.
  
      Args:
        feature_spec: A feature_spec for the input data.
        input_data_dict: Dict[str, List[Dict[str, primitive]]] the input data used
          for analysis.
        preprocessing_fn: The preprocessing_fn used for analysis.
        cache_dict: Dict[str, Dict[str, List[bytes]]], input cache dict. If
          provided, should_read_cache must be False.
        should_read_cache: A bool indicating if the pipeline should read cache. If
          True, cache_dict must be False.
        datasets_to_transform: List[str], list of dataset keys to transform.
        expected_transform_data: List[Dict[str, primitive]], the expected
          transformed data, should be the same for each dataset.
        expected_cache: Dict[str, Dict[str, bytes]], expected encoded cache.
        transform_fn_output_dir: A directory where the output transform_fn should
          be written to, if None provided it will not be written.
        use_tf_compat_v1: If True, TFT's public APIs (e.g. AnalyzeDataset) will
          use Tensorflow in compat.v1 mode. Defaults to `True`.
  
      Returns:
        A _RunPipelineResult.
      """
      input_metadata = (
          dataset_metadata.DatasetMetadata.from_feature_spec(
              feature_spec
          )
      )
      with self._TestPipeline() as p:
          with tft_beam.Context(force_tf_compat_v1=use_tf_compat_v1):
              # Wraps each value in input_data_dict as a PCollection.
              input_data_pcoll_dict = {}
              for a, b in input_data_dict.items():
                  pcoll = p | a.key >> beam.Create(b)
                  input_data_pcoll_dict[a] = pcoll
  
              pcoll_cache_dict = {}
  
              # If provided with a cache dictionary this wraps cache entries in
              # PCollections.
              if cache_dict is not None:
                  assert not should_read_cache
                  for dataset in cache_dict:
                      cache_entry = {}
                      for idx, (k, v) in enumerate(
                          cache_dict[dataset].items()
                      ):
                          cache_entry[k] = (
                              p
                              | f"CreateCache[{dataset}][{idx}]"
                              >> beam.Create(v)
                          )
                      metadata = (
                          p
                          | f"CreateCacheMetadata[{dataset}]"
                          >> beam.Create([cache_dict[dataset].metadata])
                      )
                      pcoll_cache_dict[
                          dataset
                      ] = analyzer_cache.DatasetCache(
                          cache_entry, metadata
                      )
  
              # If requested, reads cache from the test cache directory.
              if should_read_cache:
                  assert cache_dict is None
                  pcoll_cache_dict = (
                      p
                      | analyzer_cache.ReadAnalysisCacheFromFS(
                          self._cache_dir, list(input_data_dict.keys())
                      )
                  )
  
              self._publish_rendered_dot_graph_file(
                  preprocessing_fn,
                  feature_spec,
                  set(input_data_dict.keys()),
                  pcoll_cache_dict,
                  use_tf_compat_v1=use_tf_compat_v1,
              )
  
              transform_fn, cache_output = (
                  input_data_pcoll_dict,
                  pcoll_cache_dict,
                  input_metadata,
              ) | "Analyze" >> tft_beam.AnalyzeDatasetWithCache(
                  preprocessing_fn
              )
              _ = (
                  cache_output
                  | "WriteCache"
                  >> analyzer_cache.WriteAnalysisCacheToFS(
                      p, self._cache_dir
                  )
              )
  
              # Transforms the requested datasets.
              if datasets_to_transform is None:
                  transformed_dataset = None
              else:
                  flattened_transform_data = [
                      input_data_pcoll_dict[d]
                      for d in datasets_to_transform
                  ] | "FlattenTransformData" >> beam.Flatten()
                  transformed_dataset = (
                      (flattened_transform_data, input_metadata),
                      transform_fn,
                  ) | "Transform" >> tft_beam.TransformDataset()
  
              # Validate the transformed data is as expected. This requires providing
              # datasets_to_transform.
              if expected_transform_data is not None:
                  assert transformed_dataset is not None
                  (
                      transformed_data,
                      unused_transformed_metadata,
                  ) = transformed_dataset
                  beam_test_util.assert_that(
                      transformed_data,
                      beam_test_util.equal_to(expected_transform_data),
                  )
  
              if expected_cache is not None:
                  for dataset in expected_cache:
                      cache_dict = cache_output[dataset].cache_dict
                      self.assertCountEqual(
                          cache_dict.keys(),
                          expected_cache[dataset].keys(),
                      )
                      beam_test_util.assert_that(
                          cache_output[dataset].metadata,
                          beam_test_util.is_not_empty(),
                          label="AssertCacheMetadata[{}]".format(
                              dataset
                          ),
                      )
                      for idx, (key, value) in enumerate(
                          expected_cache[dataset].items()
                      ):
                          beam_test_util.assert_that(
                              cache_dict[key],
                              beam_test_util.equal_to(value),
                              label="AssertCache[{}][{}]".format(
                                  dataset, idx
                              ),
                          )
  
              # Write transform_fn if provided with an output directory.
              tft_output = None
              if transform_fn_output_dir is not None:
                  if not transform_fn_output_dir:
                      transform_fn_output_dir = os.path.join(
                          self.base_test_dir, uuid.uuid4().hex
                      )
                  _ = transform_fn | tft_beam.WriteTransformFn(
                      transform_fn_output_dir
                  )
                  tft_output = tft.TFTransformOutput(
                      transform_fn_output_dir
                  )
  
      return _RunPipelineResult(cache_output, p.metrics, tft_output)
  
imports: import functools
  import os
  import struct
  from typing import Callable, Mapping, List
  import uuid
  import apache_beam
  from apache_beam.testing import util
  import numpy
  import tensorflow
  import tensorflow_transform
  from tensorflow_transform import analyzer_nodes
  from tensorflow_transform import common_types
  from tensorflow_transform import impl_helper
  from tensorflow_transform import nodes
  import tensorflow_transform.beam
  from tensorflow_transform.beam import analysis_graph_builder
  from tensorflow_transform.beam import analyzer_cache
  from tensorflow_transform.beam import tft_unit
  from tensorflow_transform.tf_metadata import dataset_metadata
  from tfx_bsl.types import tfx_namedtuple
label_desc: Use image node_id 6 for calling the CachedImplTest obj's underlying member method code with example usage: obj._run_pipeline(feature_spec, input_data_dict, preprocessing_fn, cache_dict, should_read_cache, datasets_to_transform, expected_transform_data, expected_cache, transform_fn_output_dir, use_tf_compat_v1) and returns: _RunPipelineResult