Introduction

TensorFlow input pipelines often start from wide tabular data where only a subset of features is needed for one model. Selecting columns early makes training faster, reduces memory use, and prevents accidental feature drift between training and inference. The tf.data API gives you a clean way to do this with deterministic transformations.

Select Feature Keys from Dictionary Elements

When each dataset element is a dictionary, the most direct pattern is to map every element to a smaller dictionary that keeps only required keys. This keeps schema decisions explicit and visible in code review.

1import tensorflow as tf
2
3features = {
4    "age": tf.constant([22.0, 35.0, 41.0, 29.0]),
5    "income": tf.constant([40.0, 82.0, 120.0, 60.0]),
6    "score": tf.constant([0.2, 0.8, 0.9, 0.6]),
7    "region_id": tf.constant([1, 2, 1, 3]),
8}
9labels = tf.constant([0, 1, 1, 0], dtype=tf.int32)
10
11ds = tf.data.Dataset.from_tensor_slices((features, labels))
12
13selected_keys = ("age", "income")
14
15def select_columns(x, y):
16    narrowed = {k: x[k] for k in selected_keys}
17    return narrowed, y
18
19train_ds = ds.map(select_columns).batch(2)
20
21for batch_x, batch_y in train_ds.take(1):
22    print("keys:", list(batch_x.keys()))
23    print("age shape:", batch_x["age"].shape)
24    print("labels:", batch_y.numpy())

Use a single selected_keys definition that both training and serving code can import. If feature choices are hard coded in several files, drift appears quickly when the schema changes.

Preserve Column Order for Dense Model Inputs

Some models expect a single dense tensor instead of a dictionary. In that case, select columns and stack them in one explicit order. The order must be fixed and documented because changing order changes model meaning.

1import tensorflow as tf
2
3features = {
4    "age": tf.constant([22.0, 35.0, 41.0, 29.0]),
5    "income": tf.constant([40.0, 82.0, 120.0, 60.0]),
6    "score": tf.constant([0.2, 0.8, 0.9, 0.6]),
7}
8labels = tf.constant([0, 1, 1, 0], dtype=tf.int32)
9
10ds = tf.data.Dataset.from_tensor_slices((features, labels))
11
12ordered_columns = ("age", "income", "score")
13
14def to_dense_matrix(x, y):
15    cols = [tf.cast(x[name], tf.float32) for name in ordered_columns]
16    matrix = tf.stack(cols, axis=-1)
17    return matrix, y
18
19matrix_ds = ds.map(to_dense_matrix).batch(2)
20
21for x_batch, y_batch in matrix_ds.take(1):
22    print("batch matrix shape:", x_batch.shape)  # (2, 3)
23    print("first row:", x_batch[0].numpy())
24    print("labels:", y_batch.numpy())

If your preprocessing team keeps feature definitions in metadata, generate ordered_columns from that source once and version it. That gives you traceability when a model is retrained months later.

Select Columns While Reading CSV Data

For CSV workflows, you can reduce work even earlier by selecting only needed columns during ingestion. This avoids reading unnecessary fields into downstream steps.

1import tensorflow as tf
2
3# Example only. Replace with a real file path.
4file_pattern = "./train.csv"
5
6selected = ["age", "income", "score", "label"]
7label_name = "label"
8
9csv_ds = tf.data.experimental.make_csv_dataset(
10    file_pattern=file_pattern,
11    batch_size=32,
12    label_name=label_name,
13    select_columns=selected,
14    num_epochs=1,
15    shuffle=False,
16)
17
18for x_batch, y_batch in csv_ds.take(1):
19    print("feature keys:", sorted(x_batch.keys()))
20    print("label dtype:", y_batch.dtype)

This pattern is useful for shared data tables where each team uses different slices of the same source. It also lowers parsing overhead in large pipelines.

Common Pitfalls

• Selecting columns late in the pipeline. Move selection close to ingestion so expensive transforms do not process unused fields.
• Inconsistent feature sets between training and inference. Store selected keys in one shared constant and reuse it in both paths.
• Unstable column ordering when producing dense tensors. Use one ordered list and treat reordering as a breaking change.
• Silent schema breakage after upstream table changes. Add a startup check that validates required keys before training starts.
• Mixing label fields into the feature map. Keep labels separate to avoid accidental leakage and wrong metrics.

Summary

• Use Dataset.map to narrow dictionary elements to the exact feature keys you need.
• For dense model input, stack selected columns in a fixed documented order.
• For CSV input, prefer select_columns at read time to reduce unnecessary parsing.
• Centralize feature selection metadata so training and serving remain aligned.
• Add lightweight schema checks to detect drift early.