TIMX 512 - row group size data migration#153
Merged
Merged
Conversation
Why these changes are being introduced: It was discovered that some parquet files, mostly those before 2025-06-27, have potentially a single, very large row group instead of maximum rows per group of 1k. This has been addressed in recent work, resulting in no files after 2025-06-28 having a single row group if total rows exceed 1k. How this addresses that need: Migration script that re-packs row groups (rewrites) when a single row group is found for a number of rows that should be split up. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-512
jonavellecuerdo
approved these changes
Jul 7, 2025
jonavellecuerdo
left a comment
jonavellecuerdo
left a comment
Contributor
There was a problem hiding this comment.
These changes make sense to me! I was curious if you can include the queries you used to generate the output in the PR description?
Contributor
Author
You bet! Almost did... and then can't remember why I didn't. Worried it might confuse the issue. Here they are: import time
from timdex_dataset_api import TIMDEXDatasetMetadata
from timdex_dataset_api.config import configure_dev_logger
configure_dev_logger()
# td = TIMDEXDataset(os.environ["TIMDEX_DATASET_LOCATION"])
# # NOTE: ergonomic mismatch... I want a TDM instance but don't need to limit...
# # could use TDM.from_location() for sure, maybe good fit?
# # but then why does it return full records? that doesn't feel right...
# td.load(current_records=True)
#################
# PRE FIX
#################
print("\n################## PRE FIX ###################")
tdm_pre = TIMDEXDatasetMetadata.from_dataset_location(
# "/Users/ghukill/dev/mit/data/timdex_dataset/prod"
"s3://ghukill-test/timdex-dataset/row-group-testing/prod"
)
tdm_pre.get_full_record_via_duckdb(
"alma:9935071191506761",
"alma-full-ingest-2025-02-28t22-33-35",
)
time.sleep(5)
#################
# POST FIX
#################
print("\n################## POST FIX ###################")
tdm_post = TIMDEXDatasetMetadata.from_dataset_location(
# "/Users/ghukill/dev/mit/data/timdex_dataset/prod_003"
"s3://ghukill-test/timdex-dataset/row-group-testing/prod_003"
)
tdm_post.get_full_record_via_duckdb(
"alma:9935071191506761",
"alma-full-ingest-2025-02-28t22-33-35",
) |
Contributor
Author
I remember why! The method It was fun and helpful for testing out a needle-in-a-haystack query, but the ergonomics don't feel right for where the method should live. Here is a sketch of that method just for fun: def get_full_record_via_duckdb(
self,
timdex_record_id: str,
run_id: str | None = None,
):
# TODO: docstring
t0 = time.perf_counter()
logger.debug(
f"Retrieving full record from dataset - timdex_record_id: {timdex_record_id}"
)
# get parquet filename
# if run_id provided, use that, else query current records and use that run_id
t1 = time.perf_counter()
if run_id:
metadata_row = (
self.conn.query(
f"""
select * from records
where timdex_record_id='{timdex_record_id}'
and run_id='{run_id}';
"""
)
.to_df()
.iloc[0]
)
else:
metadata_row = (
self.conn.query(
f"""
select * from current_records
where timdex_record_id='{timdex_record_id}'
"""
)
.to_df()
.iloc[0]
)
logger.debug(
f"Parquet file identified: {metadata_row.filename}, "
f"run_record_offset: {metadata_row.run_record_offset}, "
f"elapsed: {time.perf_counter() - t1}"
)
# retrieve full record from parquet file
t1 = time.perf_counter()
query = f"""
select
timdex_record_id,
decode(source_record)::text as source_record,
decode(transformed_record)::json as transformed_record,
source,
run_date,
run_type,
action,
run_id,
run_record_offset,
run_timestamp
from read_parquet('{metadata_row.filename}')
where timdex_record_id = '{metadata_row.timdex_record_id}'
and run_id = '{metadata_row.run_id}'
and run_record_offset = {metadata_row.run_record_offset}
;
"""
full_record = self.conn.query(query).to_df().iloc[0]
logger.debug(
f"Full record retrieved, elapsed: {time.perf_counter() - t1}, "
f"total elapsed: {time.perf_counter()-t0}"
)
return full_record |
This file contains hidden or 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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Purpose and background context
It was discovered that some parquet files, mostly those before
2025-06-28, may have a single, large row group instead of multiple row groups with a maximum row count of 1k (when the parquet file has 1k+ total records).This has been addressed in recent work, resulting in no files on or after
2025-06-28having a single row group if total rows exceed 1k.This migration script re-packs row groups (rewrites) in a parquet file when a single row group is found for a number of rows that should be split up.
How can a reviewer manually see the effects of these changes?
A clone of production, and a clone with the migration applied, has been written to a testing location:
s3://ghukill-test/timdex-dataset/row-group-testing/.You can point the script with the
--dry-runflag to see that the clone still shows files that need updating, while the other dataset with the script applies skips all files because they have correct row groups.1- Set AWS dev credentials
2- Run dry run for pre-fix dataset:
3- Run dry run for post-fix dataset:
While difficult to recreate, I've also performed some "needle in a haystack" queries to retrieve records using their
run_record_offsetto help prune row groups. The dataset with the fix show significant -- orders of magnitude -- faster retrieval because libraries like DuckDB and pyarrow can now skip row groups where the sortedrun_record_offsetwon't be found.Here is some output from these tests; happy to discuss more or demonstrate how to recreate if people are curious:
The key thing to look for is the time
"Full record retrieved, elapsed"in the last line of each, which demonstrates the speedup. We go from ~11 seconds to ~0.8 seconds, which is almost pure network I/O of downloading bytes we didn't need.It's worth noting that generating the dataset metadata is a bit slower now due to more row groups. But as discussed, we are willing to pay a slight time penalty for generating the metadata, in the name of speeding up other operations.
Includes new or updated dependencies?
NO
Changes expectations for external applications?
NO
What are the relevant tickets?