code: def test_write_tfrecords(
      with_tf_schema,
      ray_start_regular_shared,
      tmp_path,
  ):
      """Test that write_tfrecords writes TFRecords correctly.
  
      Test this by writing a Dataset to a TFRecord (function under test),
      reading it back out into a tf.train.Example,
      and checking that the result is analogous to the original Dataset.
      """
  
      import tensorflow as tf
  
      # The dataset we will write to a .tfrecords file.
      ds = ray.data.from_items(
          data_partial(with_tf_schema),
          # Here, we specify `parallelism=1` to ensure that all rows end up in the same
          # block, which is required for type inference involving
          # partially missing columns.
          parallelism=1,
      )
  
      # The corresponding tf.train.Example that we would expect to read
      # from this dataset.
      expected_records = tf_records_partial()
  
      tf_schema = None
      if with_tf_schema:
          features = expected_records[0].features
          tf_schema = _features_to_schema(features)
  
      # Perform the test.
      # Write the dataset to a .tfrecords file.
      ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
  
      # Read the Examples back out from the .tfrecords file.
      # This follows the offical TFRecords tutorial:
      # https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2
  
      filenames = sorted(os.listdir(tmp_path))
      filepaths = [
          os.path.join(tmp_path, filename) for filename in filenames
      ]
      raw_dataset = tf.data.TFRecordDataset(filepaths)
  
      tfrecords = []
      for raw_record in raw_dataset:
          example = tf.train.Example()
          example.ParseFromString(raw_record.numpy())
          tfrecords.append(example)
  
      assert tfrecords == expected_records
  
imports: import json
  import os
  from typing import TYPE_CHECKING
  import numpy
  import pytest
  from pandas.api.types import is_float_dtype, is_int64_dtype, is_object_dtype
  import ray
  from ray.tests.conftest import *
label_desc: Use image node_id 9 for calling a global function with example usage: test_write_tfrecords(with_tf_schema, ray_start_regular_shared, tmp_path) without return types