Merge pull request #2939 from mabel-dev/#2927

Time travel Iceberg data

Author: joocer

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1889
+__build__ = 1893
 __author__ = "@joocer"
-__version__ = "0.26.2-beta.1889"
+__version__ = "0.26.2-beta.1893"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

opteryx/connectors/aws_s3_connector.py

@@ -20,7 +20,7 @@
 from opteryx.connectors.base.base_connector import BaseConnector
 from opteryx.connectors.capabilities import Asynchronous
 from opteryx.connectors.capabilities import Cacheable
-from opteryx.connectors.capabilities import Partitionable
+from opteryx.connectors.capabilities import Diachronic
 from opteryx.connectors.capabilities import PredicatePushable
 from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DataError
@@ -36,7 +36,7 @@
 
 
 class AwsS3Connector(
-    BaseConnector, Cacheable, Partitionable, PredicatePushable, Asynchronous, Statistics
+    BaseConnector, Cacheable, Diachronic, PredicatePushable, Asynchronous, Statistics
 ):
     __mode__ = "Blob"
     __type__ = "S3"
@@ -67,7 +67,7 @@ def __init__(self, credentials=None, **kwargs):
             raise MissingDependencyError(err.name) from err
 
         BaseConnector.__init__(self, **kwargs)
-        Partitionable.__init__(self, **kwargs)
+        Diachronic.__init__(self, **kwargs)
         Cacheable.__init__(self, **kwargs)
         PredicatePushable.__init__(self, **kwargs)
         Asynchronous.__init__(self, **kwargs)

opteryx/connectors/capabilities/__init__.py

@@ -5,16 +5,16 @@
 
 from opteryx.connectors.capabilities.asynchronous import Asynchronous
 from opteryx.connectors.capabilities.cacheable import Cacheable
+from opteryx.connectors.capabilities.diachronic import Diachronic
 from opteryx.connectors.capabilities.limit_pushable import LimitPushable
-from opteryx.connectors.capabilities.partitionable import Partitionable
 from opteryx.connectors.capabilities.predicate_pushable import PredicatePushable
 from opteryx.connectors.capabilities.statistics import Statistics
 
 __all__ = (
     "Asynchronous",
     "Cacheable",
     "LimitPushable",
-    "Partitionable",
+    "Diachronic",
     "PredicatePushable",
     "Statistics",
 )

…connectors/capabilities/partitionable.py …yx/connectors/capabilities/diachronic.pyopteryx/connectors/capabilities/partitionable.py renamed to opteryx/connectors/capabilities/diachronic.py opteryx/connectors/capabilities/partitionable.py renamed to opteryx/connectors/capabilities/diachronic.py

@@ -7,7 +7,14 @@
 from opteryx.exceptions import InvalidConfigurationError
 
 
-class Partitionable:
+class Diachronic:
+    """Capability for connectors that support diachronic (time-travel) reads.
+
+    Historically this capability was named `Partitionable`; it stores a partition
+    scheme and optional start/end date attributes used by connectors to support
+    date-range reads and time-travel queries.
+    """
+
     partitioned = True
 
     def __init__(self, **kwargs):

opteryx/connectors/disk_connector.py

@@ -23,8 +23,8 @@
 from orso.types import OrsoTypes
 
 from opteryx.connectors.base.base_connector import BaseConnector
+from opteryx.connectors.capabilities import Diachronic
 from opteryx.connectors.capabilities import LimitPushable
-from opteryx.connectors.capabilities import Partitionable
 from opteryx.connectors.capabilities import PredicatePushable
 from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DataError
@@ -37,7 +37,7 @@
 OS_SEP = os.sep
 
 
-class DiskConnector(BaseConnector, Partitionable, PredicatePushable, LimitPushable, Statistics):
+class DiskConnector(BaseConnector, Diachronic, PredicatePushable, LimitPushable, Statistics):
     """
     Connector for reading datasets from files on local storage.
     """
@@ -75,7 +75,7 @@ def __init__(self, **kwargs):
                 Arbitrary keyword arguments.
         """
         BaseConnector.__init__(self, **kwargs)
-        Partitionable.__init__(self, **kwargs)
+        Diachronic.__init__(self, **kwargs)
         PredicatePushable.__init__(self, **kwargs)
         LimitPushable.__init__(self, **kwargs)
         Statistics.__init__(self, **kwargs)

opteryx/connectors/gcp_cloudstorage_connector.py

@@ -16,7 +16,7 @@
 from opteryx.connectors.base.base_connector import BaseConnector
 from opteryx.connectors.capabilities import Asynchronous
 from opteryx.connectors.capabilities import Cacheable
-from opteryx.connectors.capabilities import Partitionable
+from opteryx.connectors.capabilities import Diachronic
 from opteryx.connectors.capabilities import PredicatePushable
 from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DatasetNotFoundError
@@ -55,7 +55,7 @@ def get_storage_credentials():
 
 
 class GcpCloudStorageConnector(
-    BaseConnector, Cacheable, Partitionable, PredicatePushable, Asynchronous, Statistics
+    BaseConnector, Cacheable, Diachronic, PredicatePushable, Asynchronous, Statistics
 ):
     __mode__ = "Blob"
     __type__ = "GCS"
@@ -89,7 +89,7 @@ def __init__(self, credentials=None, **kwargs):
             raise MissingDependencyError(name) from err
 
         BaseConnector.__init__(self, **kwargs)
-        Partitionable.__init__(self, **kwargs)
+        Diachronic.__init__(self, **kwargs)
         Cacheable.__init__(self, **kwargs)
         PredicatePushable.__init__(self, **kwargs)
         Asynchronous.__init__(self, **kwargs)

opteryx/connectors/iceberg_connector.py

@@ -20,13 +20,15 @@
 from orso.tools import single_item_cache
 from orso.types import OrsoTypes
 
-from opteryx.connectors import DiskConnector
 from opteryx.connectors.base.base_connector import BaseConnector
+from opteryx.connectors.capabilities import Diachronic
 from opteryx.connectors.capabilities import LimitPushable
 from opteryx.connectors.capabilities import PredicatePushable
 from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DatasetNotFoundError
+from opteryx.exceptions import DatasetReadError
 from opteryx.exceptions import NotSupportedError
+from opteryx.exceptions import UnsupportedSyntaxError
 from opteryx.managers.expression import NodeType
 from opteryx.managers.expression import get_all_nodes_of_type
 from opteryx.models import RelationStatistics
@@ -117,7 +119,7 @@ def _predicate_to_iceberg_filter(root):
     return iceberg_filter if iceberg_filter else "True", unsupported
 
 
-class IcebergConnector(BaseConnector, LimitPushable, Statistics, PredicatePushable):
+class IcebergConnector(BaseConnector, Diachronic, LimitPushable, Statistics, PredicatePushable):
     __mode__ = "Blob"
     __type__ = "ICEBERG"
 
@@ -140,22 +142,55 @@ class IcebergConnector(BaseConnector, LimitPushable, Statistics, PredicatePushab
         OrsoTypes.DATE,
     }
 
-    def __init__(self, *args, catalog=None, io=DiskConnector, **kwargs):
+    def __init__(self, *args, catalog=None, **kwargs):
         BaseConnector.__init__(self, **kwargs)
         LimitPushable.__init__(self, **kwargs)
+        Diachronic.__init__(self, **kwargs)
         Statistics.__init__(self, **kwargs)
         PredicatePushable.__init__(self, **kwargs)
 
         import pyiceberg
 
         try:
             self.table = catalog.load_table(self.dataset)
-            self.io_connector = io(**kwargs)
+            self.snapshot = self.table.current_snapshot()
+            self.snapshot_id = self.snapshot.snapshot_id
         except pyiceberg.exceptions.NoSuchTableError:
             raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__) from None
 
     def get_dataset_schema(self) -> RelationSchema:
-        iceberg_schema = self.table.schema()
+        if self.start_date != self.end_date:
+            if self.start_date.date() != self.end_date.date():
+                raise UnsupportedSyntaxError("This table only supports point in time reads.")
+            raise UnsupportedSyntaxError(
+                "This table only supports point in time reads. Are you missing the time component from your FOR clause?"
+            )
+
+        if self.start_date is not None:
+            snapshots = self.table.inspect.snapshots().sort_by("committed_at")
+            snapshot_rows = snapshots.to_pylist()
+
+            if not snapshot_rows:
+                raise DatasetReadError("No data available for the specified date.")
+
+            # Honor dates before the first snapshot, reject dates beyond the newest snapshot
+            if self.start_date < snapshot_rows[0]["committed_at"]:
+                selected = snapshot_rows[0]
+            elif self.start_date > snapshot_rows[-1]["committed_at"]:
+                raise DatasetReadError("No data available for the specified date.")
+            else:
+                selected = snapshot_rows[0]
+                for candidate in snapshot_rows:
+                    if candidate["committed_at"] <= self.start_date:
+                        self.statistics.dataset_committed_at = candidate["committed_at"].isoformat()
+                        selected = candidate
+                    else:
+                        break
+
+            self.snapshot_id = selected["snapshot_id"]
+            self.snapshot = self.table.snapshot_by_id(self.snapshot_id)
+
+        iceberg_schema = self.table.schemas()[self.snapshot.schema_id]
         arrow_schema = iceberg_schema.as_arrow()
 
         self.schema = RelationSchema(
@@ -169,7 +204,7 @@ def get_dataset_schema(self) -> RelationSchema:
         column_names = {col.field_id: col.name for col in iceberg_schema.columns}
         column_types = {col.field_id: col.field_type for col in iceberg_schema.columns}
 
-        files = self.table.inspect.files()
+        files = self.table.inspect.files(snapshot_id=self.snapshot_id)
 
         # No files = empty table, no stats
         if len(files.column("file_path")) == 0:
@@ -225,7 +260,10 @@ def read_dataset(
         )
 
         reader = self.table.scan(
-            row_filter=pushed_filters, selected_fields=selected_columns, limit=limit
+            row_filter=pushed_filters,
+            selected_fields=selected_columns,
+            limit=limit,
+            snapshot_id=self.snapshot_id,
         ).to_arrow_batch_reader()
 
         batch = None

opteryx/connectors/virtual_data.py

@@ -20,7 +20,7 @@
 
 from opteryx.connectors.base.base_connector import BaseConnector
 from opteryx.connectors.base.base_connector import DatasetReader
-from opteryx.connectors.capabilities import Partitionable
+from opteryx.connectors.capabilities import Diachronic
 from opteryx.connectors.capabilities import Statistics
 from opteryx.exceptions import DatasetNotFoundError
 from opteryx.utils import arrow
@@ -66,13 +66,13 @@ def suggest(dataset):
         )
 
 
-class SampleDataConnector(BaseConnector, Partitionable, Statistics):
+class SampleDataConnector(BaseConnector, Diachronic, Statistics):
     __mode__ = "Internal"
     __type__ = "SAMPLE"
 
     def __init__(self, *args, **kwargs):
         BaseConnector.__init__(self, **kwargs)
-        Partitionable.__init__(self, **kwargs)
+        Diachronic.__init__(self, **kwargs)
         Statistics.__init__(self, **kwargs)
         self.dataset = self.dataset.lower()
         self.variables = None

opteryx/planner/binder/binder_visitor.py

@@ -947,7 +947,7 @@ def visit_scan(self, node: Node, context: BindingContext) -> Tuple[Node, Binding
         from opteryx.connectors import connector_factory
         from opteryx.connectors.capabilities import Asynchronous
         from opteryx.connectors.capabilities import Cacheable
-        from opteryx.connectors.capabilities import Partitionable
+        from opteryx.connectors.capabilities import Diachronic
         from opteryx.connectors.capabilities import Statistics
         from opteryx.connectors.capabilities.cacheable import async_read_thru_cache
         from opteryx.connectors.capabilities.cacheable import read_thru_cache
@@ -969,7 +969,7 @@ def visit_scan(self, node: Node, context: BindingContext) -> Tuple[Node, Binding
 
         if hasattr(node.connector, "variables"):
             node.connector.variables = context.connection.variables
-        if Partitionable in connector_capabilities:
+        if Diachronic in connector_capabilities:
             node.connector.start_date = node.start_date
             node.connector.end_date = node.end_date
         if Cacheable in connector_capabilities and "NO_CACHE" not in (node.hints or []):
@@ -1031,7 +1031,7 @@ def visit_subquery(self, node: Node, context: BindingContext) -> Tuple[Node, Bin
         node, context = self.visit_exit(node, context)
 
         # Extract the column names to check for duplicates
-        column_names = (n.schema_column.name for n in node.columns)
+        column_names = (n.current_name for n in node.columns)
         seen = set()
         duplicates = [name for name in column_names if name in seen or seen.add(name)]  # type: ignore
 

opteryx/virtual_datasets/planet_data.py

@@ -77,8 +77,8 @@ def read(end_date=None, *args) -> pyarrow.Table:
         # April 26, 1781 - Uranus discovered by Sir William Herschel
         return full_set.take([0, 1, 2, 3, 4, 5])
     if end_date < datetime.datetime(1846, 11, 13):
-        # November 13, 1846 - Neptune
-        return full_set.take([0, 1, 2, 3, 4, 5, 7])
+        # November 13, 1846 - Neptune discovered, so only planets through Uranus exist
+        return full_set.take([0, 1, 2, 3, 4, 5, 6])
     if end_date < datetime.datetime(1930, 3, 13):
         # March 13, 1930 - Pluto discovered by Clyde William Tombaugh
         return full_set.take([0, 1, 2, 3, 4, 5, 6, 7])