migrate tests and docs snippets to yield in transformations

Author: sh-rp

docs/website/docs/general-usage/transformations/index.md

@@ -90,7 +90,7 @@ Most of the following examples will be using the ibis expressions of the `dlt.Da
 
 * **Decorator arguments** mirror those accepted by `@dlt.resource`.
 * The transformation function signature must contain at least one `dlt.Dataset` which is used inside the function to create the transformation SQL statements and calculate the resulting schema update.
-* Return a `TReadableRelation` created with ibis expressions or a select query which will be materialized into the destination table. _Do **not** yield Python dictionaries._
+* Yields a `TReadableRelation` created with ibis expressions or a select query which will be materialized into the destination table. If the first item yielded is a valid sql query or relation object, data will be interpreted as a transformation. In all other cases, the tranformation decorator will work like any other resource.
 
 ## Loading to other datasets
 

docs/website/docs/general-usage/transformations/transformation-snippets.py

@@ -36,7 +36,7 @@ def basic_transformation_snippet(fruitshop_pipeline: dlt.Pipeline) -> None:
     def copied_customers(dataset: dlt.Dataset) -> Any:
         # Ibis expression: sort by name and keep first 5 rows
         customers_table = dataset["customers"]
-        return customers_table.order_by("name").limit(5)
+        yield customers_table.order_by("name").limit(5)
 
     # Same pipeline & same dataset
     fruitshop_pipeline.run(copied_customers(fruitshop_pipeline.dataset()))
@@ -55,7 +55,7 @@ def orders_per_user_snippet(fruitshop_pipeline: dlt.Pipeline) -> None:
     @dlt.transformation(name="orders_per_user", write_disposition="merge")
     def orders_per_user(dataset: dlt.Dataset) -> Any:
         purchases = dataset["purchases"]
-        return purchases.group_by(purchases.customer_id).aggregate(order_count=purchases.id.count())
+        yield purchases.group_by(purchases.customer_id).aggregate(order_count=purchases.id.count())
 
     # @@@DLT_SNIPPET_END orders_per_user
     fruitshop_pipeline.run(orders_per_user(fruitshop_pipeline.dataset()))
@@ -69,7 +69,7 @@ def loading_to_other_datasets_snippet(fruitshop_pipeline: dlt.Pipeline) -> None:
     @dlt.transformation()
     def copied_customers(dataset: dlt.Dataset) -> Any:
         customers_table = dataset["customers"]
-        return customers_table.order_by(customers_table.name).limit(5)
+        yield customers_table.order_by(customers_table.name).limit(5)
 
     # Same duckdb instance, different dataset
     dest_p = dlt.pipeline(
@@ -98,12 +98,12 @@ def my_transformations(dataset: dlt.Dataset) -> Any:
         def enriched_purchases(dataset: dlt.Dataset) -> Any:
             purchases = dataset["purchases"]
             customers = dataset["customers"]
-            return purchases.join(customers, purchases.customer_id == customers.id)
+            yield purchases.join(customers, purchases.customer_id == customers.id)
 
         @dlt.transformation(write_disposition="replace")
         def total_items_sold(dataset: dlt.Dataset) -> Any:
             purchases = dataset["purchases"]
-            return purchases.aggregate(total_qty=purchases.quantity.sum())
+            yield purchases.aggregate(total_qty=purchases.quantity.sum())
 
         return enriched_purchases(dataset), total_items_sold(dataset)
 
@@ -129,7 +129,7 @@ def sql_queries_snippet(fruitshop_pipeline: dlt.Pipeline) -> None:
     @dlt.transformation()
     def copied_customers(dataset: dlt.Dataset) -> Any:
         customers_table = dataset("SELECT * FROM customers LIMIT 5 ORDER BY name")
-        return customers_table
+        yield customers_table
 
     # @@@DLT_SNIPPET_END sql_queries_short
 
@@ -140,7 +140,7 @@ def enriched_purchases(dataset: dlt.Dataset) -> Any:
             "SELECT customers.name, purchases.quantity FROM purchases JOIN customers ON"
             " purchases.customer_id = customers.id"
         )
-        return enriched_purchases
+        yield enriched_purchases
 
     # You can even use a different dialect than the one used by the destination by supplying the dialect parameter
     # dlt will compile the query to the right destination dialect
@@ -151,7 +151,7 @@ def enriched_purchases_postgres(dataset: dlt.Dataset) -> Any:
             " purchases.customer_id = customers.id",
             query_dialect="duckdb",
         )
-        return enriched_purchases
+        yield enriched_purchases
 
     # @@@DLT_SNIPPET_END sql_queries
 
@@ -227,7 +227,7 @@ def enriched_purchases(dataset: dlt.Dataset) -> Any:
             "SELECT customers.name, purchases.quantity FROM purchases JOIN customers ON"
             " purchases.customer_id = customers.id"
         )
-        return enriched_purchases
+        yield enriched_purchases
 
     # Let's run the transformation and see that the name column in the NEW table is also marked as PII
     fruitshop_pipeline.run(enriched_purchases(fruitshop_pipeline.dataset()))
@@ -272,7 +272,7 @@ def in_transit_transformations_snippet() -> None:
     def orders_per_store(dataset: dlt.Dataset) -> Any:
         orders = dataset["orders"]
         stores = dataset["stores"]
-        return (
+        yield (
             orders.join(stores, orders.store_id == stores.id)
             .group_by(stores.name)
             .aggregate(order_count=orders.id.count())
@@ -312,9 +312,7 @@ def cleaned_customers(dataset: dlt.Dataset) -> Any:
         customers_table = dataset.customers
 
         # filter only new customers and exclude the name column in the result
-        return customers_table.filter(customers_table.id > max_pimary_key).drop(
-            customers_table.name
-        )
+        yield customers_table.filter(customers_table.id > max_pimary_key).drop(customers_table.name)
 
     # create a warehouse dataset, would ordinarily be snowflake or some other warehousing destination
     warehouse_pipeline = dlt.pipeline(

tests/load/test_configuration.py

@@ -8,7 +8,7 @@
 def test_transformation_defaults() -> None:
     @dlt.transformation()
     def my_tf(dataset: SupportsReadableDataset[Any]) -> Any:
-        return dataset["example_table"].limit(5)
+        yield dataset["example_table"].limit(5)
 
     assert my_tf.write_disposition == "append"
     # assert my_tf(dataset).materialization == "table"

tests/load/transformations/test_basic_transformations.py

@@ -37,7 +37,7 @@ def test_simple_query_transformations(destination_config: DestinationTestConfigu
 
     @dlt.transformation()
     def copied_customers(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(5)
+        yield dataset["customers"].limit(5)
 
     # transform into transformed dataset
     dest_p.run(copied_customers(fruit_p.dataset()))
@@ -68,11 +68,11 @@ def test_grouped_transformations(destination_config: DestinationTestConfiguratio
 
     @dlt.transformation()
     def copied_customers(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(5)
+        yield dataset["customers"].limit(5)
 
     @dlt.transformation()
     def copied_customers2(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(7)
+        yield dataset["customers"].limit(7)
 
     @dlt.source()
     def transformations(dataset: dlt.Dataset) -> List[Any]:
@@ -113,7 +113,7 @@ def test_replace_sql_transformations(destination_config: DestinationTestConfigur
 
     @dlt.transformation(write_disposition="replace")
     def copied_customers(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(5)
+        yield dataset["customers"].limit(5)
 
     # transform into same dataset
     dest_p.run(copied_customers(fruit_p.dataset()))
@@ -128,7 +128,7 @@ def copied_customers(dataset: dlt.Dataset) -> Any:
         table_name="copied_customers",
     )
     def copied_customers_updated(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(3)
+        yield dataset["customers"].limit(3)
 
     # transform into same dataset
     dest_p.run(copied_customers_updated(fruit_p.dataset()))
@@ -151,7 +151,7 @@ def test_append_sql_transformations(destination_config: DestinationTestConfigura
 
     @dlt.transformation(write_disposition="append")
     def copied_customers(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(5)
+        yield dataset["customers"].limit(5)
 
     # transform into same dataset
     dest_p.run(copied_customers(fruit_p.dataset()))
@@ -162,7 +162,7 @@ def copied_customers(dataset: dlt.Dataset) -> Any:
 
     @dlt.transformation(write_disposition="append", table_name="copied_customers")
     def copied_table_updated(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].limit(7)
+        yield dataset["customers"].limit(7)
 
     # transform into same dataset
     dest_p.run(copied_table_updated(fruit_p.dataset()))
@@ -194,15 +194,15 @@ def test_sql_transformation_with_unknown_column_types(
 
     @dlt.transformation()
     def mutated_purchases(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].mutate(new_col=5).limit(5)
+        yield dataset["customers"].mutate(new_col=5).limit(5)
 
     # problem should already be detected at extraction time
     with pytest.raises(PipelineStepFailed):
         dest_p.extract(mutated_purchases(fruit_p.dataset()))
 
     @dlt.transformation()
     def mutated_purchases_with_hints(dataset: dlt.Dataset) -> Any:
-        return dataset["customers"].mutate(new_col=5).limit(5)
+        yield dataset["customers"].mutate(new_col=5).limit(5)
 
     dest_p.run(mutated_purchases_with_hints(fruit_p.dataset()))
     assert load_table_counts(dest_p, "mutated_purchases_with_hints") == {

tests/load/transformations/test_incremental_transforms.py

@@ -97,7 +97,7 @@ def transformed_items(dataset: dlt.Dataset, last_loaded_load_id: str) -> Any:
         dlt.current.resource_state()[LAST_PROCESSED_LOAD_ID] = max_load_id
 
         # return filtered transformation
-        return items_table.filter(
+        yield items_table.filter(
             items_table._dlt_load_id > last_processed_load_id,
             items_table._dlt_load_id <= last_loaded_load_id,
         ).mutate(double_items=items_table.id * 2)
@@ -149,7 +149,7 @@ def transformed_items(dataset: dlt.Dataset) -> Any:
 
         # return filtered transformation
         items_table = dataset.items
-        return items_table.filter(items_table.id > max_pimary_key).mutate(
+        yield items_table.filter(items_table.id > max_pimary_key).mutate(
             double_items=items_table.id * 2
         )
 
@@ -203,7 +203,7 @@ def transformed_items(dataset: dlt.Dataset) -> Any:
                 items_table._dlt_load_id <= last_loaded_load_id,
             ).mutate(double_items=items_table.id * 2)
 
-        return transformed_items(dataset)
+        yield transformed_items(dataset)
 
     # first round
     inc_p.run(first_load())
@@ -242,7 +242,7 @@ def test_merge_based_incremental_transform(
     def transformed_items(dataset: dlt.Dataset) -> Any:
         # return filtered transformation
         items_table = dataset.items
-        return items_table.mutate(double_items=items_table.id * 2)
+        yield items_table.mutate(double_items=items_table.id * 2)
 
     # first round
     inc_p.run(first_load())

tests/load/transformations/test_multidataset_transformations.py

@@ -9,7 +9,7 @@
 def test_combine_two_datasets(fruit_p: dlt.Pipeline, private_fruit_p: dlt.Pipeline) -> None:
     @dlt.transformation()
     def customers_with_ages(dataset: dlt.Dataset, dataset2: dlt.Dataset) -> Any:
-        return dataset["customers"].join(
+        yield dataset["customers"].join(
             dataset2["customers_ages"], dataset["customers"].id == dataset2["customers_ages"].id
         )
 

tests/load/transformations/test_transformation_lineage.py

@@ -36,7 +36,7 @@ def test_simple_lineage(
     def enriched_purchases(dataset: dlt.Dataset) -> Any:
         purchases = dataset["purchases"]
         customers = dataset["customers"]
-        return purchases.join(customers, purchases.customer_id == customers.id)
+        yield purchases.join(customers, purchases.customer_id == customers.id)
 
     dest_p.run(enriched_purchases(fruit_p.dataset()))