How about adding a feature to pass the key when performing map on DatasetDict?
CONTRIBUTOR
@lhoestq 
If it's okay with you, can I work on this?
2025-01-13T07:51:20Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2586418754
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aKZ5C
https://api.github.com/repos/huggingface/datasets/issues/comments/2586418754/reactions
2025-01-13T07:51:20Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2586418754
https://avatars.githubusercontent.com/u/93233241?v=4
https://api.github.com/users/jp1924/events{/privacy}
https://api.github.com/users/jp1924/followers
https://api.github.com/users/jp1924/following{/other_user}
https://api.github.com/users/jp1924/gists{/gist_id}
https://github.com/jp1924
jp1924
U_kgDOBY6gWQ
https://api.github.com/users/jp1924/orgs
https://api.github.com/users/jp1924/received_events
https://api.github.com/users/jp1924/repos
https://api.github.com/users/jp1924/starred{/owner}{/repo}
https://api.github.com/users/jp1924/subscriptions
User
https://api.github.com/users/jp1924
public
MEMBER
Hi ! Can you give an example of what it would look like to use this new feature ?

Note that currently you can already do

```python
ds["train"] = ds["train"].map(process_train)
ds["test"] = ds["test"].map(process_test)
```
2025-01-13T11:30:01Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2586855791
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aMElv
https://api.github.com/repos/huggingface/datasets/issues/comments/2586855791/reactions
2025-01-13T11:30:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2586855791
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
@lhoestq 
Thanks for the response! 
Let me clarify what I'm looking for with an example:

Currently, we need to write separate processing functions or call .map() separately:
```python
# Current approach
def process_train(example):
    # Training-specific processing
    return example

def process_valid(example):
    # Validation-specific processing
    return example

ds["train"] = ds["train"].map(process_train)
ds["valid"] = ds["valid"].map(process_valid)
```

What I'm proposing is to have a single processing function that knows which split it's processing:

```python
# Proposed feature
def process(example, split_key):
    if split_key == "train":
        # Training-specific processing
    elif split_key == "valid":
        # Validation-specific processing
    return example

# Using with_key=True to pass the split information
ds = ds.map(process, with_key=True)
```

This becomes particularly useful when:
1. The processing logic is heavily shared between splits but needs minor adjustments
2. You want to maintain the processing logic in one place for better maintainability
3. The processing function is complex and you want to avoid duplicating code

So I wanted to request this feature to achieve this kind of functionality. 
I've created a draft PR implementing this: https://github.com/huggingface/datasets/pull/7240/files

2025-01-13T12:23:57Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2586966711
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aMfq3
https://api.github.com/repos/huggingface/datasets/issues/comments/2586966711/reactions
2025-01-13T12:31:17Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2586966711
https://avatars.githubusercontent.com/u/93233241?v=4
https://api.github.com/users/jp1924/events{/privacy}
https://api.github.com/users/jp1924/followers
https://api.github.com/users/jp1924/following{/other_user}
https://api.github.com/users/jp1924/gists{/gist_id}
https://github.com/jp1924
jp1924
U_kgDOBY6gWQ
https://api.github.com/users/jp1924/orgs
https://api.github.com/users/jp1924/received_events
https://api.github.com/users/jp1924/repos
https://api.github.com/users/jp1924/starred{/owner}{/repo}
https://api.github.com/users/jp1924/subscriptions
User
https://api.github.com/users/jp1924
public
MEMBER
I see ! I think it makes sense, and it's more readable than doing something like this:
```python
from functools import partial
ds = DatasetDict({key: ds[key].map(partial(process, split_key=key)) for key in ds})
```

PS: you named the argument `with_key` in your example, but it might be even clearer with it's named `with_split` maybe no ?
2025-01-13T13:26:50Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2587103014
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aNA8m
https://api.github.com/repos/huggingface/datasets/issues/comments/2587103014/reactions
2025-01-13T13:26:50Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2587103014
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
CONTRIBUTOR
@lhoestq I agree. 
It seems better to use `with_split`.
So can I open a PR with this change?
2025-01-13T14:16:51Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2587223730
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aNeay
https://api.github.com/repos/huggingface/datasets/issues/comments/2587223730/reactions
2025-01-13T14:24:11Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2587223730
https://avatars.githubusercontent.com/u/93233241?v=4
https://api.github.com/users/jp1924/events{/privacy}
https://api.github.com/users/jp1924/followers
https://api.github.com/users/jp1924/following{/other_user}
https://api.github.com/users/jp1924/gists{/gist_id}
https://github.com/jp1924
jp1924
U_kgDOBY6gWQ
https://api.github.com/users/jp1924/orgs
https://api.github.com/users/jp1924/received_events
https://api.github.com/users/jp1924/repos
https://api.github.com/users/jp1924/starred{/owner}{/repo}
https://api.github.com/users/jp1924/subscriptions
User
https://api.github.com/users/jp1924
public
MEMBER
Sure !
2025-01-13T14:30:46Z
https://github.com/huggingface/datasets/issues/7356#issuecomment-2587259622
https://api.github.com/repos/huggingface/datasets/issues/7356
IC_kwDODunzps6aNnLm
https://api.github.com/repos/huggingface/datasets/issues/comments/2587259622/reactions
2025-01-13T14:30:46Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2587259622
https://avatars.githubusercontent.com/u/42851186?v=4
https://api.github.com/users/lhoestq/events{/privacy}
https://api.github.com/users/lhoestq/followers
https://api.github.com/users/lhoestq/following{/other_user}
https://api.github.com/users/lhoestq/gists{/gist_id}
https://github.com/lhoestq
lhoestq
MDQ6VXNlcjQyODUxMTg2
https://api.github.com/users/lhoestq/orgs
https://api.github.com/users/lhoestq/received_events
https://api.github.com/users/lhoestq/repos
https://api.github.com/users/lhoestq/starred{/owner}{/repo}
https://api.github.com/users/lhoestq/subscriptions
User
https://api.github.com/users/lhoestq
public
### Feature request

Add a feature to pass the key of the DatasetDict when performing map

### Motivation

I often preprocess using map on DatasetDict. 
Sometimes, I need to preprocess train and valid data differently depending on the task. 
So, I thought it would be nice to pass the key (like train, valid) when performing map on DatasetDict. 

What do you think?

### Your contribution

I can submit a pull request to add the feature to pass the key of the DatasetDict when performing map.