tensorflow读取hdfs文件（parquet、csv）

tensorflow API读取csv

官网地址：https://tensorflow.google.cn/api_docs/python/tf/data/experimental/make_csv_dataset?hl=en

tf.data.experimental.make_csv_dataset(
    file_pattern,
    batch_size,
    column_names=None,
    column_defaults=None,
    label_name=None,
    select_columns=None,
    field_delim=',',
    use_quote_delim=True,
    na_value='',
    header=True,
    num_epochs=None,
    shuffle=True,
    shuffle_buffer_size=10000,
    shuffle_seed=None,
    prefetch_buffer_size=None,
    num_parallel_reads=None,
    sloppy=False,
    num_rows_for_inference=100,
    compression_type=None,
    ignore_errors=False
)

args释义：

file_pattern  包含 CSV 记录的文件或文件路径模式的列表。请参阅tf.io.gfile.glob模式规则。
batch_size  一个 int 表示要在单个批次中组合的记录数。
column_names  按顺序对应于 CSV 列的可选字符串列表。输入记录的每列一个。如果未提供，则从记录的第一行推断列名。这些名称将是每个数据集元素的特征字典的键。
column_defaults CSV 字段的默认值的可选列表。输入记录的每个选定列一个项目。列表中的每个项目要么是有效的 CSV 数据类型（float32、float64、int32、int64 或字符串），要么是 Tensor上述类型之一。张量可以是标量默认值（如果列是可选的），也可以是空张量（如果需要列）。如果提供了 dtype 而不是张量，则该列也将根据需要进行处理。如果未提供此列表，则尝试根据读取指定文件的前 num_rows_for_inference 行来推断类型，并假定所有列都是可选的，默认0 为数值和""字符串值。如果同时 指定了 this 和 ，则它们必须具有相同的长度，并且 select_columnscolumn_defaults假定按列索引递增的顺序排序。
label_name  与标签列对应的可选字符串。如果提供，此列的数据将作为Tensor与特征字典分开的数据返回，以便数据集符合or输入函数所期望的格式。 tf.Estimator.traintf.Estimator.evaluate
select_columns  整数索引或字符串列名的可选列表，指定要选择的 CSV 数据列的子集。如果提供了列名，则这些必须对应于 在文件标题行中提供或推断的名称。指定此参数时，将仅解析和返回与指定列相对应的 CSV 列的子集。使用它会导致更快的解析和更低的内存使用。如果同时指定了 this 和 ，则它们必须具有相同的长度，并且假定按列索引递增的顺序进行排序。 column_namescolumn_defaultscolumn_defaults
field_delim 一个可选的string. 默认为. 用于分隔记录中的字段的字符分隔符。 ","
use_quote_delim 一个可选的布尔值。默认为True. 如果为 false，则将双引号视为字符串字段内的常规字符。
na_value  识别为 NA/NaN 的附加字符串。
header  一个布尔值，指示提供的 CSV 文件的第一行是否对应于具有列名的标题行，并且不应包含在数据中。
num_epochs  指定此数据集重复次数的 int。如果没有，则永远循环遍历数据集。
shuffle 一个布尔值，指示输入是否应该被打乱。
shuffle_buffer_size 用于洗牌的缓冲区大小。较大的缓冲区大小可确保更好的洗牌，但会增加内存使用和启动时间。
shuffle_seed  用于洗牌的随机化种子。
prefetch_buffer_size  一个 int 指定要预取以提高性能的特征批次的数量。推荐值是每个训练步骤消耗的批次数。默认为自动调谐。
num_parallel_reads  用于从文件中读取 CSV 记录的线程数。如果>1，结果将被交错。默认为1.
sloppy  如果True，将以非确定性排序为代价提高读取性能。如果False，则产生的元素的顺序在洗牌之前是确定的（如果 ，元素仍然是随机shuffle=True的。请注意，如果设置了种子，则洗牌后的元素顺序是确定的）。默认为False.
num_rows_for_inference  如果未提供 record_defaults，则用于类型推断的文件行数。如果为 None，则读取所有文件的所有行。默认为 100。
compression_type  （可选。）计算为（无压缩）、 或tf.string之一的标量 。默认为无压缩。 """ZLIB""GZIP"
ignore_errors （可选。）如果True，忽略 CSV 文件解析错误，例如格式错误的数据或空行，并继续下一个有效的 CSV 记录。否则，数据集会在遇到任何无效记录时引发错误并停止处理。默认为False.

eg：

CSV_COLUMNS = ['survived', 'sex', 'age', 'n_siblings_spouses', 'parch', 'fare', 'class', 'deck', 'embark_town', 'alone']
TRAIN_DATA_URL = "hdfs:///user/hive/warehouse/xxx/dt=20220808/*/train/*.csv"

dataset = tf.data.experimental.make_csv_dataset(
     file_pattern=TRAIN_DATA_URL,
     batch_size=10,
     label_name=LABEL_COLUMN,
     column_names=CSV_COLUMNS,
     shuffle=True,
     shuffle_buffer_size=10000,
     num_epochs=None,
     ...)

tensorflow 读取parquet

线上的任务目前都是存储在数仓，为了节省资源也为了任务效率，都是snappy+parquet格式存储在HDFS ,所以为了方便读取，根据make_csv_dataset改写了下：

def make_parquet_dataset(
    file_pattern: Union[str, List[str]],
    batch_size: int = None,
    select_columns: List[str] = None,
    shuffle: bool = False,
    shuffle_buffer_size: int = 100000,
    shuffle_seed: int = None,
    num_epochs: int = 1,
    prefetch_buffer_size: int = None,
    block_length: int = 512,
    n_limit_file: int = None,
    dataset_name: str = "",
    col2spec: Dict = None,
) -> tf.data.Dataset:
    """
    参照 tf.data.make_csv_dataset()改写

    Args:
      file_pattern: 输入数据集路径模式, 格式比如"hdfs:///xxx/dt=20220227/*"
      batch_size:
      selected_columns: 选择的列名列表, 其他列被丢弃. 默认所有列
      shuffle: 是否shuffle, 如果true, 速度会变慢
      shuffle_buffer_size: in sample, not batch
      shuffle_seed:
      num_epochs: 数据集轮数，如果大于1, dataset会repeat
      prefetch_buffer_size: 如果batch_size>0, 是batch; 否则, 是sample
      block_length: >1, faster

    Returns:
      dataset
    """
    # FIXME: 频繁一直报警告: WARN hdfs.DFSClient: zero
    if isinstance(file_pattern, str):
        f_list = tf.io.gfile.glob(file_pattern)
    elif isinstance(file_pattern, list):
        f_list = []
        for pattern in file_pattern:
            for f in tf.io.gfile.glob(file_pattern):
                f_list.append(f)
    else:
        raise ValueError("unsupported file_pattern")

    f_list = sorted(f_list)

    if n_limit_file:
        f_list = f_list[:n_limit_file]

    # print(f"{dataset_name}: {f_list}")
    dataset = tf.data.Dataset.from_tensor_slices(f_list) \
                             .shuffle(len(f_list), seed=shuffle_seed)

    if col2spec is None:  # 自动推断col2spec, not work in distributed mode
        sample = tfio.IODataset.from_parquet(f_list[0]).take(1)
        col2spec = dict(sample.element_spec)

    if select_columns is not None:
        col2spec = {col: spec for (col, spec) in col2spec.items() if col.decode() in select_columns}

    def filename_to_dataset(f_name: str):
        dataset = tfio.IODataset.from_parquet(f_name, columns=col2spec)
        return dataset

    dataset = dataset.interleave(
        map_func=filename_to_dataset,
        # cycle_length=len(f_list),
        cycle_length=4, #4个线程读文件
        block_length=block_length,
        num_parallel_calls=tf.data.AUTOTUNE,
        deterministic=True,
    )


"""
shuffle/repeat和repeat/shuffle的顺序对最终结果有一些影响
"""
    if shuffle:  
        dataset = dataset.shuffle(shuffle_buffer_size, seed=shuffle_seed)

    if num_epochs != 1:
        dataset = dataset.repeat(num_epochs)

    if batch_size is not None:
        dataset = dataset.batch(batch_size=batch_size,
                                drop_remainder=True,
                                )
    if prefetch_buffer_size is not None:
        dataset = dataset.prefetch(prefetch_buffer_size)

    return dataset

用法：

make_parquet_dataset(
                            file_pattern="hdfs:///xx/xxx/*",
                            batch_size=1000
                            prefetch_buffer_size=10,
                            shuffle=True,
                            shuffle_buffer_size=10000,
                            dataset_name="train",
                            col2spec=['name','age']
                        )

• shuffle(buffsize) 用于将数据打乱，其中buffsize的大小越大，数据的混乱程度越高，因为shuffle的实现思路为：开辟一可容纳buffsize个数据的缓冲区，初始时将数据的前buffsize个读入缓冲区，而后随机在缓冲区里选择一个输出，同时将数据的第buffsize+1个读入缓冲区。
• repeat(count) 用于将数据重复count次，相当于我们训练时的epoch。
• batch(batch_size)用于将数据划分为多个batch，同时tensorflow中有着很好的调整功能，当最后一个batch不满足batchsize时就以当前长度输出。
• repeat后的数据进行打乱，这会使得不同epcoh间的数据被打乱，即前一个epcoh中数据未加载完，下一个epoch中数据可能插入，导致一个epoch中可能数据重复多次；
• 先shuffle再repeat，epoch内部打乱，一定先输出完一个epoch内所有值；