code: def create_tf_record_dataset(
      filenames,
      parse_function,
      batch_size,
      file_parallelism=None,
      shuffle_buffer_size=0,
  ):
      dataset = tf.data.TFRecordDataset(
          filenames, num_parallel_reads=file_parallelism
      )
      if shuffle_buffer_size > 0:
          dataset = dataset.shuffle(shuffle_buffer_size)
      dataset = dataset.batch(batch_size)
      dataset = dataset.map(parse_function)
      dataset = dataset.prefetch(1)
      return dataset
  
file_path: io/tests/test_atds_avro/utils/benchmark_utils.py
imports: import os
  import tensorflow
  from tests.test_atds_avro.utils.data_source_registry import get_canonical_name, get_data_source_from_registry
  from tests.test_atds_avro.utils.generator.tensor_generator import IntTensorGenerator, FloatTensorGenerator, BoolTensorGenerator
  from tests.test_atds_avro.utils.generator.sparse_tensor_generator import IntSparseTensorGenerator, ValueDistribution
  from tests.test_atds_avro.utils.generator.varlen_tensor_generator import WordVarLenTensorGenerator, DimensionDistribution
  from tests.test_atds_avro.utils.tf_record_writer import TFRecordWriter
label_desc: Use image node_id 2 for calling a global function with example usage: create_tf_record_dataset(filenames, parse_function, batch_size, file_parallelism, shuffle_buffer_size) and returns: dataset