Tensorflow
Dataset API
input pipeline
parquet files
machine learning

Tensorflow Dataset API input pipeline with parquet files

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Building a tf.data pipeline from Parquet files is mostly about choosing the ingestion layer, because core TensorFlow does not behave like it does for TFRecord where a native reader is the default path. In practice, the two common approaches are reading Parquet batches with pyarrow and feeding them into tf.data.Dataset.from_generator, or using TensorFlow I/O when its Parquet support matches your environment.

Why Parquet Needs An Extra Step

Parquet is a columnar analytics format. That makes it excellent for storage and offline processing, but it is not the canonical TensorFlow training format. TensorFlow's most direct integrations are built around tensors, arrays, and record-based formats.

So the design question is not whether Parquet is valid. It is how to bridge columnar files into batches of tensors efficiently.

A good baseline is:

  • read Parquet in batches
  • convert each batch into NumPy arrays or Python scalars
  • emit dictionaries that map feature names to tensors
  • let tf.data handle shuffling, batching, and prefetching after that point

A Runnable pyarrow Plus tf.data Pipeline

The following example creates a small Parquet file, reads it in record batches, and turns it into a Dataset.

python
1from pathlib import Path
2import numpy as np
3import pyarrow as pa
4import pyarrow.parquet as pq
5import tensorflow as tf
6
7path = Path("train.parquet")
8
9table = pa.table({
10    "age": pa.array([21, 34, 19, 42], type=pa.int32()),
11    "income": pa.array([50_000.0, 88_000.0, 32_000.0, 120_000.0], type=pa.float32()),
12    "label": pa.array([0, 1, 0, 1], type=pa.int32()),
13})
14pq.write_table(table, path)
15
16def parquet_rows(file_path):
17    parquet_file = pq.ParquetFile(file_path)
18    for batch in parquet_file.iter_batches(batch_size=2):
19        rows = batch.to_pydict()
20        for i in range(len(rows["label"])):
21            yield {
22                "age": np.int32(rows["age"][i]),
23                "income": np.float32(rows["income"][i]),
24                "label": np.int32(rows["label"][i]),
25            }
26
27output_signature = {
28    "age": tf.TensorSpec(shape=(), dtype=tf.int32),
29    "income": tf.TensorSpec(shape=(), dtype=tf.float32),
30    "label": tf.TensorSpec(shape=(), dtype=tf.int32),
31}
32
33dataset = tf.data.Dataset.from_generator(
34    lambda: parquet_rows(str(path)),
35    output_signature=output_signature,
36)
37
38dataset = dataset.shuffle(4).batch(2).prefetch(tf.data.AUTOTUNE)
39
40for batch in dataset.take(1):
41    print(batch)

This is simple and reliable because pyarrow already understands Parquet well and tf.data only has to reason about tensors after the conversion boundary.

Splitting Features And Labels

Most training code wants features and labels separated. You can do that with a map step.

python
1def split_features_and_label(row):
2    label = row.pop("label")
3    return row, label
4
5train_ds = dataset.map(split_features_and_label)
6
7for features, label in train_ds.take(1):
8    print(features)
9    print(label)

That keeps the ingestion layer clean and lets you add normalization or feature engineering later.

When TensorFlow I/O Helps

TensorFlow I/O provides additional filesystem and format integrations, including Parquet-related APIs in supported setups. If you are already using TensorFlow I/O successfully in your environment, it can reduce the amount of custom generator code you write.

That said, generator-based ingestion with pyarrow is often easier to debug because:

  • schema inspection is explicit
  • conversion to NumPy types is under your control
  • version mismatches are easier to isolate

If a training pipeline must be boring and dependable, explicit conversion is often a strength rather than a drawback.

Performance Considerations

The slow part of a Parquet pipeline is usually not the Dataset object. It is conversion and I/O.

To keep throughput healthy:

  • read in batches rather than one Parquet row group cell at a time
  • avoid Python object-heavy transformations once data is in tf.data
  • batch after ingestion and use prefetch
  • keep dtypes aligned with model expectations so TensorFlow does not keep casting tensors later

If you have many files, interleave them instead of reading one full file at a time. A practical pattern is one generator per file with an outer dataset of file names.

Schema Discipline Matters

Parquet supports rich schemas, optional fields, and nested structures. That flexibility is useful for storage, but training pipelines benefit from stricter expectations.

Before training, make sure:

  • column names are stable
  • nullability is handled explicitly
  • categorical fields are converted intentionally
  • every emitted feature has a fixed TensorFlow dtype

A pipeline that sometimes emits Python int, sometimes NumPy int64, and sometimes missing values will become fragile quickly.

Why Not Just Convert To TFRecord

For large, repeated training jobs, converting Parquet to TFRecord ahead of time is often worth considering. TFRecord integrates more naturally with TensorFlow and can simplify distributed training setups.

Parquet still makes sense when:

  • it is your source-of-truth format
  • the same files feed Spark, Pandas, and TensorFlow
  • you want schema tooling outside TensorFlow

So the decision is operational, not ideological.

Common Pitfalls

  • Assuming TensorFlow core has the same first-class Parquet reader story as TFRecord.
  • Letting from_generator emit inconsistent Python types across rows.
  • Ignoring null values until tensor conversion fails late in the pipeline.
  • Building a row-by-row Python pipeline with no batching and then blaming tf.data for poor throughput.
  • Treating Parquet schema flexibility as if it were harmless to model input contracts.

Summary

  • A practical TensorFlow Parquet pipeline usually reads with pyarrow and feeds rows or batches into tf.data.
  • 'Dataset.from_generator is a dependable bridge when Parquet support is not native enough for your case.'
  • TensorFlow I/O can be useful, but explicit conversion is often easier to debug.
  • Keep dtypes, null handling, and feature names consistent from the start.
  • Consider TFRecord if training performance and TensorFlow-native tooling matter more than keeping Parquet end to end.

Course illustration
Course illustration

All Rights Reserved.