diff --git a/src/sempy_labs/_notebooks.py b/src/sempy_labs/_notebooks.py
index 0538ac9e..63a179c0 100644
--- a/src/sempy_labs/_notebooks.py
+++ b/src/sempy_labs/_notebooks.py
@@ -163,7 +163,7 @@ def import_notebook_from_web(
     if len(dfI_filt) == 0:
         create_notebook(
             name=notebook_name,
-            notebook_content=response.content.decode('utf-8'),
+            notebook_content=response.content.decode("utf-8"),
             workspace=workspace_id,
             description=description,
             format="ipynb",
@@ -216,7 +216,9 @@ def create_notebook(
         Defaults to None which places the notebook in the root of the workspace.
     """
 
-    notebook_payload = base64.b64encode(notebook_content.encode('utf-8')).decode("utf-8")
+    notebook_payload = base64.b64encode(notebook_content.encode("utf-8")).decode(
+        "utf-8"
+    )
 
     definition_payload = {
         "parts": [
diff --git a/src/sempy_labs/directlake/_update_directlake_partition_entity.py b/src/sempy_labs/directlake/_update_directlake_partition_entity.py
index 867cad1c..46b71f4d 100644
--- a/src/sempy_labs/directlake/_update_directlake_partition_entity.py
+++ b/src/sempy_labs/directlake/_update_directlake_partition_entity.py
@@ -12,6 +12,7 @@
 from typing import List, Optional, Union
 import sempy_labs._icons as icons
 from uuid import UUID
+import json
 
 
 @log
@@ -76,6 +77,7 @@ def update_direct_lake_partition_entity(
                 for p in t.Partitions
                 if t.Name == tName
             )
+            current_slt = tom.model.Tables[tName].SourceLineageTag
 
             if part_name is None:
                 raise ValueError(
@@ -85,14 +87,56 @@ def update_direct_lake_partition_entity(
             tom.model.Tables[tName].Partitions[part_name].Source.EntityName = eName
 
             # Update source lineage tag
-            existing_schema = (
-                tom.model.Tables[tName].Partitions[part_name].Source.SchemaName or "dbo"
-            )
-            if schema is None:
-                schema = existing_schema
+            if schema:
+                # Only set schema for DL over SQL (not DL over OneLake)
+                expression_source_name = (
+                    tom.model.Tables[tName]
+                    .Partitions[part_name]
+                    .Source.ExpressionSource.Name
+                )
+                expr = tom.model.Expressions[expression_source_name].Expression
+                if "Sql.Database" in expr:
+                    tom.model.Tables[tName].Partitions[
+                        part_name
+                    ].Source.SchemaName = schema
+                tom.model.Tables[tName].SourceLineageTag = f"[{schema}].[{eName}]"
+            else:
+                tom.model.Tables[tName].SourceLineageTag = f"[dbo].[{eName}]"
+
+            new_slt = tom.model.Tables[tName].SourceLineageTag
+
+            # PBI_RemovedChildren logic
+            try:
+                e_name = (
+                    tom.model.Tables[tName]
+                    .Partitions[part_name]
+                    .Source.ExpressionSource.Name
+                )
+                ann = tom.get_annotation_value(
+                    object=tom.model.Expressions[e_name], name="PBI_RemovedChildren"
+                )
+                if ann:
+                    e = json.loads(ann)
+                    for i in e:
+                        sltag = (
+                            i.get("remoteItemId", {})
+                            .get("analysisServicesObject", {})
+                            .get("sourceLineageTag", {})
+                        )
+                        if sltag == current_slt:
+                            i["remoteItemId"]["analysisServicesObject"][
+                                "sourceLineageTag"
+                            ] = new_slt
+                    tom.set_annotation(
+                        object=tom.model.Expressions[e_name],
+                        name="PBI_RemovedChildren",
+                        value=json.dumps(e),
+                    )
+            except Exception as e:
+                print(
+                    f"{icons.yellow_dot} Warning: Could not update PBI_RemovedChildren annotation for table '{tName}'. {str(e)}"
+                )
 
-            tom.model.Tables[tName].Partitions[part_name].Source.SchemaName = schema
-            tom.model.Tables[tName].SourceLineageTag = f"[{schema}].[{eName}]"
             print(
                 f"{icons.green_dot} The '{tName}' table in the '{dataset_name}' semantic model within the '{workspace_name}' workspace has been updated to point to the '{eName}' table."
             )