Tensorflow Dataset API input pipeline with parquet files
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.
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.datahandle 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.
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.
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_generatoremit 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.datafor poor throughput. - Treating Parquet schema flexibility as if it were harmless to model input contracts.
Summary
- A practical TensorFlow Parquet pipeline usually reads with
pyarrowand feeds rows or batches intotf.data. - '
Dataset.from_generatoris 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
TFRecordif training performance and TensorFlow-native tooling matter more than keeping Parquet end to end.
Related reading
- Tensorflow Dataset API not using GPU
- Tensorflow Dataset API with HDFS
- Tensorflow dataset data preprocessing is done once for the whole dataset or for each call to iterator.next?
- TensorFlow DataSet from_generator with variable batch size
- TensorFlow Dataset Generator With Mixed Datatypes
- Tensorflow Dataset .map API
- Tensorflow Estimator API Summaries
- Tensorflow Layers Api Linear Activation Function

System Design Fundamentals
Build a strong foundation in designing scalable, reliable distributed systems.
View the courseTrack what you have practised
A free account saves your progress, solutions and study plan across every problem on Codemia.
ML System Design practice on Codemia
Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.