Skip to content

[fix] handle parquet schema mismatch in dataset concatenation #146

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

Merged
kcz358 merged 2 commits into from
Mar 19, 2026
Merged

[fix] handle parquet schema mismatch in dataset concatenation #146

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c8cc421
[fix] handle parquet schema mismatch in dataset concatenation
mwxely Mar 18, 2026
e28d773
refactor: use cast() instead of to_list() for schema alignment
mwxely Mar 19, 2026
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
32 changes: 30 additions & 2 deletions src/lmms_engine/utils/data_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,34 @@
FPS_MAX_FRAMES = 768


def _safe_concatenate_datasets(data_list):
"""Concatenate datasets with automatic schema alignment.

When loading multiple parquet files, schema mismatches can occur if columns
have different inferred types (e.g., a struct column is null in one file but
has nested fields in another). This function handles such cases by falling
back to row-wise concatenation when direct concatenation fails.

Args:
data_list: List of HuggingFace Dataset objects to concatenate.

Returns:
A single concatenated Dataset.
"""
if len(data_list) <= 1:
return concatenate_datasets(data_list)
try:
return concatenate_datasets(data_list)
except Exception as e:
logger.warning(
f"Direct concatenation failed due to schema mismatch: {e}. " f"Falling back to row-wise concatenation."
)
all_rows = []
for ds in data_list:
all_rows.extend(ds.to_list())
return Dataset.from_list(all_rows)
Copy link
Copy Markdown
Collaborator

@kcz358 kcz358 Mar 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be cause extremely slow and large memory needed under large scale data. Recommend to preprocess or cast dataset before training to ensure the dataset is in the same features.

Copy link
Copy Markdown
Collaborator Author

@mwxely mwxely Mar 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point — I've replaced the to_list() fallback with cast()-based schema alignment in e28d773. The new approach:

  1. Tries cast() to the first dataset's features (fast, zero-copy for compatible schemas)
  2. If that fails, builds a merged feature set preferring non-null types and casts all datasets to it

This keeps the fix zero-overhead on the happy path and avoids materializing the full dataset into memory on schema mismatch.



class DataUtilities:
@staticmethod
def load_json(path: str) -> List[Dict[str, List]]:
Expand Down Expand Up @@ -104,7 +132,7 @@ def load_yaml(path: str) -> Tuple[List[Dict[str, List]], List[str]]:
data_list.append(data)
logger.info(f"Dataset size: {len(data)}")
data_folder_list.extend([data_folder] * len(data))
data_list = concatenate_datasets(data_list)
data_list = _safe_concatenate_datasets(data_list)
return data_list, data_folder_list

@staticmethod
Expand Down Expand Up @@ -222,7 +250,7 @@ def load_inline_datasets(
data_list.append(data)
logger.info(f"Dataset size: {len(data)}")
data_folder_list.extend([data_folder] * len(data))
data_list = concatenate_datasets(data_list)
data_list = _safe_concatenate_datasets(data_list)

return data_list, data_folder_list

Expand Down
Loading