Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Only keep columns in produces #490

Merged
merged 7 commits into from
Oct 11, 2023
Merged

Only keep columns in produces #490

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
88dec2c
Only keep columns in produces
PhilippeMoussalli Oct 5, 2023
3c5fa19
fix ordering
PhilippeMoussalli Oct 5, 2023
eec8f0a
implement PR feedback
PhilippeMoussalli Oct 10, 2023
ede7b2b
Merge branch 'main' into select-columns-load-component
PhilippeMoussalli Oct 10, 2023
7a00a2b
address PR feedback
PhilippeMoussalli Oct 10, 2023
2858023
Merge branch 'main' into select-columns-load-component
PhilippeMoussalli Oct 10, 2023
71fd45a
Merge branch 'main' into select-columns-load-component
RobbeSneyders Oct 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions components/load_from_hf_hub/src/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,20 @@ def __init__(self,
def load(self) -> dd.DataFrame:
# 1) Load data, read as Dask dataframe
logger.info("Loading dataset from the hub...")
dask_df = dd.read_parquet(f"hf://datasets/{self.dataset_name}")

# Only read required columns
columns = []

invert_column_name_mapping = {v: k for k, v in self.column_name_mapping.items()}
for subset_name, subset in self.spec.produces.items():
for field_name, field in subset.fields.items():
subset_field_name = f"{subset_name}_{field_name}"
column_name = invert_column_name_mapping.get \
(subset_field_name, subset_field_name)
columns.append(column_name)

logger.debug(f"Columns to keep: {columns}")
dask_df = dd.read_parquet(f"hf://datasets/{self.dataset_name}", columns=columns)

# 2) Make sure images are bytes instead of dicts
if self.image_column_names is not None:
Expand All @@ -71,7 +84,7 @@ def load(self) -> dd.DataFrame:
dask_df = dask_df.head(self.n_rows_to_load, npartitions=npartitions)
dask_df = dd.from_pandas(dask_df, npartitions=npartitions)

# 4) Set the index
# 5) Set the index
if self.index_column == "None":
logger.info(
"Index column not specified, setting a globally unique index",
Expand Down
15 changes: 14 additions & 1 deletion components/load_from_parquet/src/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,20 @@ def __init__(self,
def load(self) -> dd.DataFrame:
# 1) Load data, read as Dask dataframe
logger.info("Loading dataset from the file...")
dask_df = dd.read_parquet(self.dataset_uri)

# Only read required columns
columns = []
if self.column_name_mapping is not None:
invert_column_name_mapping = {v: k for k, v in self.column_name_mapping.items()}
for subset_name, subset in self.spec.produces.items():
for field_name, field in subset.fields.items():
subset_field_name = f"{subset_name}_{field_name}"
column_name = invert_column_name_mapping.get \
(subset_field_name, subset_field_name)
columns.append(column_name)

logger.debug(f"Columns to keep: {columns}")
dask_df = dd.read_parquet(self.dataset_uri, columns=columns)

# 2) Rename columns
if self.column_name_mapping:
Expand Down
Loading