Use `tf.data` to efficiently count the number of Examples in TFRecords?

I'm writing some simple code to count the number of Examples in TFRecord files. Previously I was using something like this:

def _count_example(path):
    return sum(1 for _ in tf.python_io.tf_record_iterator(path))

def _count_total_examples(tfrecord_glob_pattern):
    example_count = 0
    paths = glob.glob(tfrecord_glob_pattern)

    with futures.ProcessPoolExecutor(max_workers=os.cpu_count()) as executor:
        results = executor.map(_count_example, paths, chunksize=max(len(paths) // os.cpu_count(), 1))
        for result in results:
            example_count += result

    return example_count

I've copied my test TFRecord input to a RAM disk to make sure we are not bottlenecked by I/O. My test TFRecord files have 635,424 image files have 128 shards. The above code finishes in 13 seconds.

I'm trying to change the same code to using tf.data and see if I can get similar results. So I'm doing something like:

def _count_total_examples(tfrecord_glob_pattern):
    dataset = tf.data.Dataset.list_files(tfrecord_glob_pattern, shuffle=False)
    dataset = tf.data.TFRecordDataset(dataset, num_parallel_reads=os.cpu_count())
    iterator = dataset.make_one_shot_iterator()
    next_element = iterator.get_next()
    with tf.Session() as sess:
        example_count = 0
        for _ in trange(635424):
            sess.run(next_element)
            example_count += 1
    return example_count

This version takes more than 2 minutes to finish. I guess this is because I'm calling sess.run() too many times and this has a non-trivial overhead. To verify this claim, I've changed the above code to:

def _count_total_examples(tfrecord_glob_pattern):
    dataset = tf.data.Dataset.list_files(tfrecord_glob_pattern, shuffle=False)
    dataset = tf.data.TFRecordDataset(dataset, num_parallel_reads=os.cpu_count())
    dataset = dataset.map(lambda _: 1, num_parallel_calls=os.cpu_count())
    dataset = dataset.batch(4096)
    iterator = dataset.make_one_shot_iterator()
    next_element = iterator.get_next()
    with tf.Session() as sess:
        example_count = 0
        for _ in trange(635424 // 4096):
            sess.run(next_element)
            example_count += 1
    return example_count

This time the code finishes within 30 seconds - which is still not ideal but I guess better. Basically, based on my above observations, I have two questions:

1) What is the efficient approach to address this issue? Or should I just stay using tf.python_io.tf_record_iterator?

2) Is there a way to "loop the whole dataset and then return the final result" in the TensorFlow graph? I've tried something like tf.contrib.data.Reducer but it seems to be still quite slow (takes around one minute to finish):

def _count_total_examples(tfrecord_glob_pattern):
    dataset = tf.data.Dataset.list_files(tfrecord_glob_pattern, shuffle=False)
    dataset = tf.data.TFRecordDataset(dataset, num_parallel_reads=os.cpu_count() * 2)
    reducer = tf.contrib.data.Reducer(
        init_func=lambda _: 0,
        reduce_func=lambda curr, _: curr + 1,
        finalize_func=lambda curr: curr
    )
    dataset = tf.contrib.data.reduce_dataset(dataset, reducer)
    with tf.Session() as sess:
        result = sess.run(dataset)
        example_count = result
    return example_count