[Data] Support class constructor args for map() and flat_map() #38606
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Signed-off-by: Cheng Su <scnju13@gmail.com>
c21
requested review from
ericl,
scv119,
amogkam,
scottjlee,
bveeramani and
raulchen
as code owners
c21
commented
Aug 18, 2023
Comment on lines
+181
to
+189
| result = ds.map_batches( | ||
| StatefulFnWithArgs, | ||
| compute=ray.data.ActorPoolStrategy(), | ||
| fn_args=(1,), | ||
| fn_kwargs={"kwarg": 2}, | ||
| fn_constructor_args=(1,), | ||
| fn_constructor_kwargs={"kwarg": 2}, | ||
| ).take() == list(range(10)) | ||
| ).take() | ||
| assert sorted(extract_values("id", result)) == list(range(10)), result |
There was a problem hiding this comment.
This is bugfix for previous test.
scottjlee
approved these changes
Aug 18, 2023
raulchen
reviewed
Aug 18, 2023
| @@ -286,6 +286,7 @@ def map( | |||
| fn: UserDefinedFunction[Dict[str, Any], Dict[str, Any]], | |||
| *, | |||
| compute: Optional[ComputeStrategy] = None, | |||
| fn_constructor_args: Optional[Iterable[Any]] = None, | |||
There was a problem hiding this comment.
should also add fn_constructor_kwargs?
There was a problem hiding this comment.
I prefer us to do the minimal we need here. We can add later if there's ask.
| raise ValueError( | ||
| "fn_constructor_args can only be specified if providing a " | ||
| f"CallableClass instance for fn, but got: {fn}" | ||
| ) |
There was a problem hiding this comment.
The type of the compute arg is too loose. This makes us need duplicated null checks.
It'd be better to split this into 2 steps.
- a
parse_computeutil function that convertsOptional[Union[str, "ComputeStrategy"]]to justComputeStrategy. - Then check the fn and args.
And by doing this, compute can only be None or str in the Dataset API level. under the hood, it's always the concrete type.
There was a problem hiding this comment.
Make sense, we already have a get_compute() method.
Let me make a separate PR to do the refactoring, it would involve change AbstractUDFMap.compute and _plan_udf_map_op.
raulchen
approved these changes
Aug 18, 2023
arvind-chandra
pushed a commit
to lmco/ray
that referenced
this pull request
Aug 31, 2023
…roject#38606) We get user request to support class constructor args for `Dataset.map()` and `Dataset.flat_map()`, similar to `Dataset.map_batches()`. No technical reason why we cannot support it, given they are got converted to same map function during execution. This PR is to add them. Signed-off-by: Cheng Su <scnju13@gmail.com> Signed-off-by: e428265 <arvind.chandramouli@lmco.com>
vymao
pushed a commit
to vymao/ray
that referenced
this pull request
Oct 11, 2023
…roject#38606) We get user request to support class constructor args for `Dataset.map()` and `Dataset.flat_map()`, similar to `Dataset.map_batches()`. No technical reason why we cannot support it, given they are got converted to same map function during execution. This PR is to add them. Signed-off-by: Cheng Su <scnju13@gmail.com> Signed-off-by: Victor <vctr.y.m@example.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Why are these changes needed?
We get user request to support class constructor args for
Dataset.map()andDataset.flat_map(), similar toDataset.map_batches(). No technical reason why we cannot support it, given they are got converted to same map function during execution. This PR is to add them.Related issue number
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.