feat(datasets): Added the Experimental PolarsDatabaseDataset #990
Conversation
Signed-off-by: Minura Punchihewa <[email protected]>
Signed-off-by: Minura Punchihewa <[email protected]>
|
Hey @noklam, @deepyaman, Should we implement Or do you have different thoughts on how this dataset ought to be implemented? |
| class PolarsDatabaseDataset(SQLQueryDataset): | ||
|
|
||
| def __init__( # noqa: PLR0913 | ||
| self, |
There was a problem hiding this comment.
| self, | |
| self, | |
| *, |
| ) | ||
|
|
||
| def save(self, data: None) -> NoReturn: | ||
| pass |
There was a problem hiding this comment.
maybe NotImplemented instead?
| def load(self) -> pl.DataFrame: | ||
| load_args = copy.deepcopy(self._load_args) | ||
|
|
||
| if self._filepath: | ||
| load_path = get_filepath_str(PurePosixPath(self._filepath), self._protocol) | ||
| with self._fs.open(load_path, mode="r") as fs_file: | ||
| query = fs_file.read() | ||
| else: | ||
| query = load_args.pop("sql") | ||
|
|
||
| return pl.read_database( | ||
| query=query, | ||
| connection=self._connection_str, | ||
| **load_args |
There was a problem hiding this comment.
Can you add some docs or checking to reflect this logic
- If filepath exist use it
- Otherwise use sql
- If both are defined, maybe error out or at least log which one is being used?
Ideally I would re-order the argument as well, since the dataset put sql as the first argument but actually filepath have higher priority which feels counter-intuitive.
Sorry for the late response; didn't see this. While there may be opportunities to reduce code copying more broadly, most datasets just inherit from
Probably, because Polars supports it. You can make the table name optional but require it for |
Description
This PR adds the
PolarsDatabaseDatasetto support interactions with databases using Polars.Fixes #853
Development notes
I have extended the
SQLQueryDatasetto implement this dataset.These changes have been tested,
1. Manually, by running the code locally to load and save tensors from and to Safetensors files.2. Via the existing and newly added unit tests.Checklist
RELEASE.mdfile