fix(tableau/ingestion): allow for changes in Glides schema (#677)

* Write failing test

* Isolate problem test case

* Pass test

* Tweak datetime regex

* Allow glides schema to be out of order

Author: runkelcorey

src/lamp_py/ingestion/glides.py

@@ -8,10 +8,7 @@
 import dataframely as dy
 import polars as pl
 import pyarrow
-import pyarrow.dataset as pd
 import pyarrow.parquet as pq
-import pyarrow.compute as pc
-from dateutil.relativedelta import relativedelta
 
 from lamp_py.aws.s3 import download_file, upload_file
 from lamp_py.aws.kinesis import KinesisReader
@@ -24,7 +21,7 @@
 )
 
 RFC3339_DATE_REGEX = r"^20(?:([1-3][0-9]-[0-1][0-9]-[0-3][0-9]))"  # up to 2039-19-39
-RFC3339_DATETIME_REGEX = RFC3339_DATE_REGEX + r"T([0-2][0-9]:[0-5][0-9]:[0-5][0-9](?:\.\d+)?)(Z|[\+-]\d{2}:\d{2})?$"
+RFC3339_DATETIME_REGEX = RFC3339_DATE_REGEX + r"[T ]([0-2][0-9]:[0-5][0-9]:[0-5][0-9](?:\.\d+)?)(Z|[\+-]\d{2}:\d{2})?$"
 GTFS_TIME_REGEX = r"^([0-9]{2}):([0-5][0-9]):([0-5][0-9])$"  # clock can be greater than 24 hours
 
 user = dy.Struct(
@@ -67,7 +64,7 @@ class GlidesRecord(dy.Schema):
     id = dy.String()
     type = dy.String()
     time = dy.Datetime(  # in %Y-%m-%dT%H:%M:%S%:z format before serialization
-        min=datetime(2024, 1, 1), max=datetime(2039, 12, 31)  # within Python's serializable range
+        min=datetime(2024, 1, 1), max=datetime(2039, 12, 31), time_unit="ms"  # within Python's serializable range
     )
     source = dy.String()
     specversion = dy.String()
@@ -213,57 +210,36 @@ def download_remote(self) -> None:
         download_file(object_path=self.remote_path, file_name=self.local_path)
 
     @abstractmethod
-    def convert_records(self) -> pd.Dataset:
+    def convert_records(self) -> dy.DataFrame[GlidesRecord]:
         """Convert incoming records into a flattened table of records"""
 
     def append_records(self) -> None:
         """Add incoming records to a local parquet file"""
         process_logger = ProcessLogger(process_name="append_glides_records", type=self.type)
         process_logger.log_start()
 
-        new_dataset = self.convert_records()
+        new_dataset = self.convert_records().lazy()
 
         if os.path.exists(self.local_path):
-            remote_records = pd.dataset(self.local_path, schema=self.get_table_schema)
-            joined_ds = pd.dataset([new_dataset, remote_records])
+            remote_records = self.table_schema.scan_parquet(self.local_path, validation="allow")
+            joined_ds = pl.union([new_dataset, remote_records])
         else:
             joined_ds = new_dataset
 
         process_logger.add_metadata(
-            new_records=new_dataset.count_rows(),
-            total_records=joined_ds.count_rows(),
+            new_records=new_dataset.select("time").count().collect().item(),
+            total_records=joined_ds.select("time").count().collect().item(),
         )
 
-        now = datetime.now()
-        start = datetime(2024, 1, 1)
-
         with tempfile.TemporaryDirectory() as tmp_dir:
 
             new_path = os.path.join(tmp_dir, self.base_filename)
-            row_group_count = 0
-            with pq.ParquetWriter(new_path, schema=self.get_table_schema) as writer:
-                while start < now:
-                    end = start + relativedelta(months=1)
-                    if end < now:
-                        row_group = pl.DataFrame(
-                            joined_ds.filter((pc.field("time") >= start) & (pc.field("time") < end)).to_table()
-                        )
-
-                    else:
-                        row_group = pl.DataFrame(joined_ds.filter((pc.field("time") >= start)).to_table())
-
-                    if not row_group.is_empty():
-                        unique_table = (
-                            row_group.unique(keep="first").sort(by=["time"]).to_arrow().cast(self.get_table_schema)
-                        )
-
-                        row_group_count += 1
-                        writer.write_table(unique_table)
-
-                    start = end
-
-            os.replace(new_path, self.local_path)
-            process_logger.add_metadata(row_group_count=row_group_count)
+            sorted_ds = joined_ds.unique().sort("time")
+            valid = process_logger.log_dataframely_filter_results(*self.table_schema.filter(sorted_ds))
+            if not valid.is_empty():
+                pq.write_table(valid.to_arrow().cast(self.get_table_schema), new_path)
+                os.replace(new_path, self.local_path)
+                process_logger.add_metadata(row_count=pq.read_metadata(self.local_path).num_rows)
 
         process_logger.log_complete()
 
@@ -290,15 +266,17 @@ def __init__(self) -> None:
     def unique_key(self) -> str:
         return "changes"
 
-    def convert_records(self) -> pd.Dataset:
+    def convert_records(self) -> dy.DataFrame[GlidesRecord]:
         process_logger = ProcessLogger(process_name="convert_records", type=self.type)
         process_logger.log_start()
 
         editors_table = pyarrow.Table.from_pylist(self.records, schema=self.get_event_schema)
         editors_table = flatten_table_schema(editors_table)
         editors_table = explode_table_column(editors_table, "data.changes")
         editors_table = flatten_table_schema(editors_table)
-        editors_dataset = pd.dataset(editors_table)
+        editors_dataset = process_logger.log_dataframely_filter_results(
+            *EditorChangesTable.filter(pl.DataFrame(editors_table))
+        )
 
         process_logger.log_complete()
         return editors_dataset
@@ -322,12 +300,14 @@ def __init__(self) -> None:
     def unique_key(self) -> str:
         return "operator"
 
-    def convert_records(self) -> pd.Dataset:
+    def convert_records(self) -> dy.DataFrame[GlidesRecord]:
         process_logger = ProcessLogger(process_name="convert_records", type=self.type)
         process_logger.log_start()
         osi_table = pyarrow.Table.from_pylist(self.records, schema=self.get_event_schema)
         osi_table = flatten_table_schema(osi_table)
-        osi_dataset = pd.dataset(osi_table)
+        osi_dataset = process_logger.log_dataframely_filter_results(
+            *OperatorSignInsTable.filter(pl.DataFrame(osi_table))
+        )
 
         process_logger.log_complete()
         return osi_dataset
@@ -348,7 +328,7 @@ def __init__(self) -> None:
     def unique_key(self) -> str:
         return "tripUpdates"
 
-    def convert_records(self) -> pd.Dataset:
+    def convert_records(self) -> dy.DataFrame[GlidesRecord]:
         def flatten_multitypes(record: Dict) -> Dict:
             """
             For each update in a record, flatten out the objects in "cars",
@@ -374,7 +354,7 @@ def flatten_multitypes(record: Dict) -> Dict:
         tu_table = flatten_table_schema(tu_table)
         tu_table = explode_table_column(tu_table, "data.tripUpdates")
         tu_table = flatten_table_schema(tu_table)
-        tu_dataset = pd.dataset(tu_table)
+        tu_dataset = process_logger.log_dataframely_filter_results(*TripUpdatesTable.filter(pl.DataFrame(tu_table)))
 
         process_logger.log_complete()
         return tu_dataset
@@ -398,13 +378,15 @@ def __init__(self) -> None:
     def unique_key(self) -> str:
         return "tripKey"
 
-    def convert_records(self) -> pd.Dataset:
+    def convert_records(self) -> dy.DataFrame[GlidesRecord]:
         process_logger = ProcessLogger(process_name="convert_records", type=self.type)
         process_logger.log_start()
 
         tu_table = pyarrow.Table.from_pylist(self.records, schema=self.get_event_schema)
         tu_table = flatten_table_schema(tu_table)
-        tu_dataset = pd.dataset(tu_table)
+        tu_dataset = process_logger.log_dataframely_filter_results(
+            *VehicleTripAssignmentTable.filter(pl.DataFrame(tu_table))
+        )
 
         process_logger.log_complete()
         return tu_dataset

src/lamp_py/tableau/jobs/glides.py

@@ -128,7 +128,7 @@ def create_trips_updated_glides_parquet(job: HyperJob, num_files: Optional[int])
                 # pl.col("data.tripUpdates.endTime").str.to_time("%H:%M:%S", strict=False),
             )
 
-            writer.write_table(polars_df.to_arrow())
+            writer.write_table(polars_df.select(job.output_processed_schema.names).to_arrow())
 
 
 def create_operator_signed_in_glides_parquet(job: HyperJob, num_files: Optional[int]) -> None:
@@ -166,7 +166,7 @@ def create_operator_signed_in_glides_parquet(job: HyperJob, num_files: Optional[
                 pl.col("time").dt.convert_time_zone(time_zone="US/Eastern").dt.replace_time_zone(None),
             )
 
-            writer.write_table(polars_df.to_arrow())
+            writer.write_table(polars_df.select(job.output_processed_schema.names).to_arrow())
 
 
 class HyperGlidesTripUpdates(HyperJob):

tests/ingestion/test_glides.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from os import remove
 from pathlib import Path
 from queue import Queue
@@ -7,7 +8,7 @@
 import dataframely as dy
 import pytest
 import polars as pl
-from polars.testing import assert_frame_equal
+import pyarrow.parquet as pq
 
 from lamp_py.ingestion.glides import (
     GlidesConverter,
@@ -37,9 +38,14 @@ def test_convert_records(dy_gen: dy.random.Generator, converter: GlidesConverter
     converter.records = converter.record_schema.sample(
         num_rows=num_rows,
         generator=dy_gen,
+        overrides={
+            "time": dy_gen.sample_datetime(
+                num_rows, min=datetime(2024, 1, 1), max=datetime(2039, 12, 31), time_unit="us"
+            ).cast(pl.Datetime(time_unit="ms"))
+        },
     ).to_dicts()
 
-    table = pl.scan_pyarrow_dataset(converter.convert_records())
+    table = converter.convert_records()
 
     assert not converter.table_schema.validate(table).is_empty()
     assert converter.table_schema.validate(table).select("id").unique().height == num_rows  # all records
@@ -48,8 +54,13 @@ def test_convert_records(dy_gen: dy.random.Generator, converter: GlidesConverter
 
 @pytest.mark.parametrize(
     ["column_transformations"],
-    [({},), ({"id": pl.col("id")},), ({"new_col": pl.lit(1)},)],
-    ids=["no-remote-records", "same-schema", "dropped-column"],
+    [
+        ({},),
+        ({"id": pl.col("id")},),
+        ({"new_col": pl.lit(1)},),
+        ({"time": pl.col("time").cast(pl.Datetime(time_unit="us")).dt.offset_by("1us")},),
+    ],
+    ids=["no-remote-records", "same-schema", "dropped-column", "truncated-timestamp"],
 )
 @pytest.mark.parametrize(
     [
@@ -67,26 +78,42 @@ def test_append_records(
     dy_gen: dy.random.Generator,
     converter: GlidesConverter,
     tmp_path: Path,
-    column_transformations: dict,
+    column_transformations: dict[str, pl.Expr],
     num_rows: int = 5,
 ) -> None:
-    """It writes all records locally."""
+    """It writes all records locally using the table schema."""
     converter.records = converter.record_schema.sample(
         num_rows=num_rows,
         generator=dy_gen,
+        overrides={
+            "time": dy_gen.sample_datetime(
+                num_rows, min=datetime(2024, 1, 1), max=datetime(2039, 12, 31), time_unit="us"
+            ).cast(pl.Datetime(time_unit="ms"))
+        },
     ).to_dicts()
 
     converter.local_path = tmp_path.joinpath(converter.base_filename).as_posix()
 
-    expectation = pl.scan_pyarrow_dataset(converter.convert_records()).collect()
+    expectation = converter.convert_records()
 
+    remote_records_height = 0
     if column_transformations:
-        remote_records = expectation.with_columns(**column_transformations)
+        remote_records = converter.table_schema.sample(
+            num_rows,
+            generator=dy_gen,
+            overrides={
+                "time": dy_gen.sample_datetime(
+                    num_rows, min=datetime(2024, 1, 1), max=datetime(2039, 12, 31), time_unit="us"
+                ).cast(pl.Datetime(time_unit="ms"))
+            },
+        ).with_columns(**column_transformations)
         remote_records.write_parquet(converter.local_path)
+        remote_records_height = remote_records.height
 
     converter.append_records()
 
-    assert_frame_equal(expectation, pl.read_parquet(converter.local_path), check_row_order=False)
+    assert pq.read_schema(converter.local_path) == converter.get_table_schema
+    assert pq.read_metadata(converter.local_path).num_rows == expectation.height + remote_records_height
 
 
 @pytest.mark.parametrize(
@@ -102,13 +129,18 @@ def test_append_records(
     ids=["editor-changes", "operator-sign-ins", "trip-updates", "vehicle-trip-assignments"],
 )
 def test_ingest_glides_events(
-    converter: GlidesConverter, dy_gen: dy.random.Generator, mocker: MockerFixture, events_per_converter: int = 500
+    converter: GlidesConverter, dy_gen: dy.random.Generator, mocker: MockerFixture, events_per_converter: int = 50
 ) -> None:
     """It routes events to correct converter and writes them to specified storage."""
     test_records = (
         converter.record_schema.sample(  # generate test records
             num_rows=events_per_converter,
             generator=dy_gen,
+            overrides={
+                "time": dy_gen.sample_datetime(
+                    events_per_converter, min=datetime(2024, 1, 1), max=datetime(2039, 12, 31), time_unit="us"
+                ).cast(pl.Datetime(time_unit="ms"))
+            },
         )
         .with_columns(
             time=pl.col("time").dt.strftime("%Y-%m-%dT%H:%M:%SZ")