Merge pull request #12 from MITLibraries/TIMX-424-reorder-partition-columns

TIMX 424 - reorder partition columns

Author: ghukill

tests/test_dataset_write.py

@@ -169,14 +169,14 @@ def test_dataset_write_schema_partitions_correctly_ordered(
             "source": "alma",
             "run_date": "2024-12-01",
             "run_type": "daily",
-            "action": "index",
             "run_id": "000-111-aaa-bbb",
+            "action": "index",
         },
     )
     file = written_files[0]
     assert (
         "/source=alma/run_date=2024-12-01/run_type=daily"
-        "/action=index/run_id=000-111-aaa-bbb" in file.path
+        "/run_id=000-111-aaa-bbb/action=index/" in file.path
     )
 
 

tests/utils.py

@@ -43,18 +43,17 @@ def generate_sample_records_with_simulated_partitions(
     run_dates = ["2024-01-01", "2024-06-15", "2024-12-31"]
     run_types = ["full", "daily"]
     actions = ["index", "delete"]
-    run_ids = [str(uuid.uuid4()) for x in range(num_run_ids)]
 
     records_remaining = num_records
     while records_remaining > 0:
         batch_size = random.randint(1, min(100, records_remaining))
+        source = random.choice(sources)
         yield from generate_sample_records(
             num_records=batch_size,
-            timdex_record_id_prefix=random.choice(sources),
-            source=random.choice(sources),
+            timdex_record_id_prefix=source,
+            source=source,
             run_date=random.choice(run_dates),
             run_type=random.choice(run_types),
             action=random.choice(actions),
-            run_id=random.choice(run_ids),
         )
         records_remaining -= batch_size

timdex_dataset_api/__init__.py

@@ -3,7 +3,7 @@
 from timdex_dataset_api.dataset import TIMDEXDataset
 from timdex_dataset_api.record import DatasetRecord
 
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = [
     "DatasetRecord",

timdex_dataset_api/dataset.py

@@ -28,17 +28,17 @@
         pa.field("source", pa.string()),
         pa.field("run_date", pa.date32()),
         pa.field("run_type", pa.string()),
-        pa.field("action", pa.string()),
         pa.field("run_id", pa.string()),
+        pa.field("action", pa.string()),
     )
 )
 
 TIMDEX_DATASET_PARTITION_COLUMNS = [
     "source",
     "run_date",
     "run_type",
-    "action",
     "run_id",
+    "action",
 ]
 
 DEFAULT_BATCH_SIZE = 1_000
@@ -54,6 +54,7 @@ def __init__(self, location: str | list[str]):
         self.dataset: ds.Dataset = None  # type: ignore[assignment]
         self.schema = TIMDEX_DATASET_SCHEMA
         self.partition_columns = TIMDEX_DATASET_PARTITION_COLUMNS
+        self._written_files: list[ds.WrittenFile] = None  # type: ignore[assignment]
 
     @classmethod
     def load(cls, location: str) -> "TIMDEXDataset":
@@ -197,6 +198,7 @@ def write(
             - use_threads: boolean if threads should be used for writing
         """
         start_time = time.perf_counter()
+        self._written_files = []
 
         if isinstance(self.source, list):
             raise TypeError(
@@ -209,14 +211,13 @@ def write(
             batch_size=batch_size,
         )
 
-        written_files = []
         ds.write_dataset(
             record_batches_iter,
             base_dir=self.source,
             basename_template="%s-{i}.parquet" % (str(uuid.uuid4())),  # noqa: UP031
             existing_data_behavior="delete_matching",
             filesystem=self.filesystem,
-            file_visitor=lambda written_file: written_files.append(written_file),
+            file_visitor=lambda written_file: self._written_files.append(written_file),  # type: ignore[arg-type]
             format="parquet",
             max_open_files=500,
             max_rows_per_file=MAX_ROWS_PER_FILE,
@@ -227,8 +228,8 @@ def write(
             use_threads=use_threads,
         )
 
-        logger.info(f"write elapsed: {round(time.perf_counter()-start_time, 2)}s")
-        return written_files  # type: ignore[return-value]
+        self.log_write_statistics(start_time)
+        return self._written_files  # type: ignore[return-value]
 
     def get_dataset_record_batches(
         self,
@@ -266,3 +267,24 @@ def get_dataset_record_batches(
                 f"elapsed: {round(time.perf_counter()-batch_start_time, 6)}s"
             )
             yield batch
+
+    def log_write_statistics(self, start_time: float) -> None:
+        """Parse written files from write and log statistics."""
+        total_time = round(time.perf_counter() - start_time, 2)
+        total_files = len(self._written_files)
+        total_rows = sum(
+            [
+                wf.metadata.num_rows  # type: ignore[attr-defined]
+                for wf in self._written_files
+            ]
+        )
+        total_size = sum(
+            [wf.size for wf in self._written_files]  # type: ignore[attr-defined]
+        )
+        logger.info(
+            f"Dataset write complete - elapsed: "
+            f"{total_time}s, "
+            f"total files: {total_files}, "
+            f"total rows: {total_rows}, "
+            f"total size: {total_size}"
+        )

timdex_dataset_api/record.py

@@ -24,8 +24,8 @@ class DatasetRecord:
     source: str | None = None
     run_date: str | datetime.datetime | None = None
     run_type: str | None = None
-    action: str | None = None
     run_id: str | None = None
+    action: str | None = None
 
     def to_dict(
         self,
@@ -46,7 +46,7 @@ def validate(self) -> None:
         # ensure all partition columns are set
         missing_partition_values = [
             field
-            for field in ["source", "run_date", "run_type", "action", "run_id"]
+            for field in ["source", "run_date", "run_type", "run_id", "action"]
             if getattr(self, field) is None
         ]
         if missing_partition_values: