ENH: Add to_tf_dataset method to convert Pandas dataframe to TensorFlow dataset

[jamiecash]

Feature Type

• Adding new functionality to pandas

• Changing existing functionality in pandas

• Removing existing functionality in pandas

Problem Description

I wish I could use pandas to create TensorFlow datasets

Feature Description

Add a new method to_tf_dataset to DataFrame, specifying columns to use as features and labels.

def to_tf_dataset(feature_columns: list[str], label_column: str) -> tf.data.Dataset:
"""
Params:
feature_columns: list[str]: A list of columns to use as features for the dataset
label_column: The column to use as the label for the dataset. Numeric columns will be used as is. Non numeric columns will be converted to pd.Categorical
"""

Alternative Solutions

Manual conversion using tf.data.Dataset.from_tensor_slices

Additional Context

No response

[jamiecash]

take

[jamiecash]

#Method

def to_tf_dataset(feature_columns : list[str], label_column: str) -> tf.data.Dataset:
    _copy = self.copy()
    
    # Check if the feature columns exist. If not, raise an KeyError exception.
    if not set(feature_columns).issubset(_copy.columns):
        raise KeyError(f"Specified feature column(s) {set(feature_columns).difference(_copy.columns)} do not exist.")
    
    # Check if the label column exists
    if label_column not in _copy.columns:
        raise KeyError(f"Specified label column '{label_column}' does not exist.")
        
    # If data type of label is not numeric, then convert to categorical and use codes as label
    if not pd.api.types.is_numeric_dtype(_copy[label_column]):
        _copy[label_column] = pd.Categorical(_copy[label_column]).codes
    
    # Pop the label
    label = _copy.pop(label_column)
    
    # Create the data source
    ds = tf.data.Dataset.from_tensor_slices((_copy[feature_columns], label))
    
    # Return it
    return ds

#Tests

• Test 1 - Invalid feature column should raise a KeyError
• Test 2 - Invalid label column should raise a KeyError
• Test 3 - Numeric label column creates dataset with numeric label
• Test 4 - Non numeric label column creates dataset with numeric label that can be converted back using pd.Categorical

[mroeschke]

Thanks for the request, but similar to #46000, I think this IO path would be too niche for pandas to specifically support and would be better as a 3rd party library.

[mroeschke]

Given that there's tf.data.Dataset.from_tensor_slices, I don't think it's sufficently necessarily to extend the pandas API to wrap this method, so closing