allow passing arbitrary transaction_rid

Author: nicornk

libs/foundry-dev-tools/src/foundry_dev_tools/resources/dataset.py

@@ -800,7 +800,7 @@ def to_polars(self) -> pl.DataFrame:
         """
         return self.query_foundry_sql("SELECT *", return_type="polars")
 
-    def to_lazy_polars(self) -> pl.LazyFrame:
+    def to_lazy_polars(self, transaction_rid: str | None = None) -> pl.LazyFrame:
         """Get dataset as a :py:class:`polars.LazyFrame`.
 
         Returns a lazy polars DataFrame that can be queried efficiently using
@@ -824,14 +824,17 @@ def to_lazy_polars(self) -> pl.LazyFrame:
         """
         from foundry_dev_tools._optional.polars import pl
 
-        last_transaction = self.get_last_transaction()
-        if last_transaction is None:
-            msg = f"Dataset has no transactions: {self.rid=}"
-            raise DatasetHasNoTransactionsError(info=msg)
+        if transaction_rid is None:
+            maybe_transaction = self.get_last_transaction()
+            if maybe_transaction is None:
+                msg = f"Dataset has no transactions: {self.rid=}"
+                raise DatasetHasNoTransactionsError(info=msg)
+            transaction_rid = maybe_transaction["rid"]
 
         return pl.scan_parquet(
-            f"s3://{self.rid}.{last_transaction['rid']}/",
+            f"s3://{self.rid}.{transaction_rid}/**/*.parquet",
             storage_options=self._context.s3.get_polars_storage_options(),
+            hive_partitioning=True,
         )
 
     @contextmanager

tests/unit/resources/test_dataset.py

@@ -12,13 +12,28 @@ def test_to_lazy_polars_no_transaction():
         ds.rid = "ri.foundry.main.dataset.test-dataset"
         ds.path = "/test/dataset/path"
 
-        # Mock get_last_transaction to return None
         with mock.patch.object(ds, "get_last_transaction", return_value=None):
-            # Assert that the correct exception is raised with the expected message
             with pytest.raises(DatasetHasNoTransactionsError) as exc_info:
                 ds.to_lazy_polars()
 
-            # Verify the error message contains the expected information
             error_message = str(exc_info.value)
             assert "Dataset has no transactions" in error_message
             assert ds.rid in error_message
+
+
+def test_to_lazy_polars_transaction_rid_logic():
+    with mock.patch.object(Dataset, "__created__", True):
+        ds = Dataset.__new__(Dataset)
+        ds.rid = "ri.foundry.main.dataset.abc123"
+        ds._context = mock.MagicMock()
+        ds._context.s3.get_polars_storage_options.return_value = {"aws_access_key_id": "test"}
+
+        with mock.patch("foundry_dev_tools._optional.polars.pl.scan_parquet") as mock_scan:
+            mock_scan.return_value = mock.MagicMock()
+            ds.to_lazy_polars(transaction_rid="test")
+
+            mock_scan.assert_called_once()
+            call_args = mock_scan.call_args
+            assert call_args[0][0] == f"s3://{ds.rid}.test/**/*.parquet"
+            assert call_args[1]["storage_options"] == ds._context.s3.get_polars_storage_options()
+            assert call_args[1]["hive_partitioning"] is True