Added Assessment step to estimate the size of DBFS Root Tables. (#741)

Closes #330 

- Added assessment step estimate dbfs root table size.
- This should help deciding whether the table is a good candidate for
cloning.

Author: FastLee

src/databricks/labs/ucx/framework/crawlers.py

@@ -43,7 +43,7 @@ def create_table(self, full_name: str, klass: Dataclass):
 
     _builtin_type_mapping: ClassVar[dict[type, str]] = {
         str: "STRING",
-        int: "INT",
+        int: "LONG",
         bool: "BOOLEAN",
         float: "FLOAT",
     }

src/databricks/labs/ucx/hive_metastore/table_size.py

@@ -0,0 +1,66 @@
+import logging
+from collections.abc import Iterable
+from dataclasses import dataclass
+from functools import partial
+
+from databricks.labs.ucx.framework.crawlers import CrawlerBase, RuntimeBackend
+from databricks.labs.ucx.hive_metastore import TablesCrawler
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TableSize:
+    catalog: str
+    database: str
+    name: str
+    size_in_bytes: int
+
+
+class TableSizeCrawler(CrawlerBase):
+    def __init__(self, backend: RuntimeBackend, schema):
+        from pyspark.sql.session import SparkSession  # type: ignore[import-not-found]
+
+        """
+        Initializes a TablesSizeCrawler instance.
+
+        Args:
+            backend (SqlBackend): The SQL Execution Backend abstraction (either REST API or Spark)
+            schema: The schema name for the inventory persistence.
+        """
+        self._backend: RuntimeBackend = backend
+        super().__init__(backend, "hive_metastore", schema, "table_size", TableSize)
+        self._tables_crawler = TablesCrawler(backend, schema)
+        self._spark = SparkSession.builder.getOrCreate()
+
+    def _crawl(self) -> Iterable[TableSize]:
+        """Crawls and lists tables using table crawler
+        Identifies DBFS root tables and calculates the size for these.
+        """
+        for table in self._tables_crawler.snapshot():
+            if not table.kind == "TABLE":
+                continue
+            if not table.is_dbfs_root():
+                continue
+            size_in_bytes = self.get_table_size(table.key)
+            yield TableSize(
+                catalog=table.catalog, database=table.database, name=table.name, size_in_bytes=size_in_bytes
+            )
+
+    def _try_load(self) -> Iterable[TableSize]:
+        """Tries to load table information from the database or throws TABLE_OR_VIEW_NOT_FOUND error"""
+        for row in self._fetch(f"SELECT * FROM {self._full_name}"):
+            yield TableSize(*row)
+
+    def snapshot(self) -> list[TableSize]:
+        """
+        Takes a snapshot of tables in the specified catalog and database.
+
+        Returns:
+            list[Table]: A list of Table objects representing the snapshot of tables.
+        """
+        return self._snapshot(partial(self._try_load), partial(self._crawl))
+
+    def get_table_size(self, table_full_name: str) -> int:
+        logger.debug(f"Evaluating {table_full_name} table size.")
+        return self._spark._jsparkSession.table(table_full_name).queryExecution().analyzed().stats().sizeInBytes()

src/databricks/labs/ucx/hive_metastore/tables.py

@@ -1,5 +1,6 @@
 import logging
 import re
+import typing
 from collections import defaultdict
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass
@@ -28,6 +29,18 @@ class Table:
 
     storage_properties: str | None = None
 
+    DBFS_ROOT_PREFIXES: typing.ClassVar[list[str]] = [
+        "/dbfs/",
+        "dbfs:/",
+    ]
+
+    DBFS_ROOT_PREFIX_EXCEPTIONS: typing.ClassVar[list[str]] = [
+        "/dbfs/mnt",
+        "dbfs:/mnt",
+        "/dbfs/databricks-datasets",
+        "dbfs:/databricks-datasets",
+    ]
+
     @property
     def is_delta(self) -> bool:
         if self.table_format is None:
@@ -80,6 +93,17 @@ def sql_unset_upgraded_to(self, catalog):
             f"UNSET TBLPROPERTIES IF EXISTS('upgraded_to');"
         )
 
+    def is_dbfs_root(self) -> bool:
+        if not self.location:
+            return False
+        for exception in self.DBFS_ROOT_PREFIX_EXCEPTIONS:
+            if self.location.startswith(exception):
+                return False
+        for prefix in self.DBFS_ROOT_PREFIXES:
+            if self.location.startswith(prefix):
+                return True
+        return False
+
 
 @dataclass
 class TableError:

src/databricks/labs/ucx/install.py

@@ -44,6 +44,7 @@
 from databricks.labs.ucx.hive_metastore.grants import Grant
 from databricks.labs.ucx.hive_metastore.hms_lineage import HiveMetastoreLineageEnabler
 from databricks.labs.ucx.hive_metastore.locations import ExternalLocation, Mount
+from databricks.labs.ucx.hive_metastore.table_size import TableSize
 from databricks.labs.ucx.hive_metastore.tables import Table, TableError
 from databricks.labs.ucx.runtime import main
 from databricks.labs.ucx.workspace_access.base import Permissions
@@ -130,6 +131,7 @@ def deploy_schema(sql_backend: SqlBackend, inventory_schema: str):
             functools.partial(table, "grants", Grant),
             functools.partial(table, "groups", MigratedGroup),
             functools.partial(table, "tables", Table),
+            functools.partial(table, "table_size", TableSize),
             functools.partial(table, "table_failures", TableError),
             functools.partial(table, "workspace_objects", WorkspaceObjectInfo),
             functools.partial(table, "permissions", Permissions),

src/databricks/labs/ucx/queries/assessment/main/01_3_all_tables.sql

@@ -1,17 +1,31 @@
 -- viz type=table, name=Table Types, search_by=name, columns=name,type,format,storage,is_delta,location
 -- widget title=Table Types, row=1, col=3, size_x=3, size_y=8
-SELECT CONCAT(`database`, '.', name) AS name,
+SELECT CONCAT(tables.`database`, '.', tables.name) AS name,
        object_type AS type,
        table_format AS format,
        CASE
            WHEN STARTSWITH(location, "dbfs:/mnt") THEN "DBFS MOUNT"
            WHEN STARTSWITH(location, "/dbfs/mnt") THEN "DBFS MOUNT"
+           WHEN STARTSWITH(location, "dbfs:/databricks-datasets") THEN "Databricks Demo Dataset"
+           WHEN STARTSWITH(location, "/dbfs/databricks-datasets") THEN "Databricks Demo Dataset"
            WHEN STARTSWITH(location, "dbfs:/") THEN "DBFS ROOT"
            WHEN STARTSWITH(location, "/dbfs/") THEN "DBFS ROOT"
            WHEN STARTSWITH(location, "wasb") THEN "UNSUPPORTED"
            WHEN STARTSWITH(location, "adl") THEN "UNSUPPORTED"
            ELSE "EXTERNAL"
        END AS storage,
        IF(format = "DELTA", "Yes", "No") AS is_delta,
-       location
-FROM $inventory.tables
+       location,
+       CASE
+            WHEN size_in_bytes IS null THEN "Non DBFS Root"
+            WHEN size_in_bytes > 10000000000000000 THEN "SIZE OUT OF RANGE"
+            WHEN size_in_bytes < 100 THEN CONCAT(CAST(size_in_bytes AS string)," Bytes")
+            WHEN size_in_bytes < 100000 THEN CONCAT(CAST(round(size_in_bytes/1024,2) AS string),"KB")
+            WHEN size_in_bytes < 100000000 THEN CONCAT(CAST(round(size_in_bytes/1024/1024,2) AS string),"MB")
+            WHEN size_in_bytes < 100000000000 THEN CONCAT(CAST(round(size_in_bytes/1024/1024/1024,2) AS string),"GB")
+            ELSE CONCAT(CAST(round(size_in_bytes/1024/1024/1024/1024,2) AS string),"TB")
+       END AS table_size
+FROM $inventory.tables left outer join $inventory.table_size on
+$inventory.tables.catalog = $inventory.table_size.catalog and
+$inventory.tables.database = $inventory.table_size.database and
+$inventory.tables.name = $inventory.table_size.name

src/databricks/labs/ucx/runtime.py

@@ -18,6 +18,7 @@
     Mounts,
     TablesCrawler,
 )
+from databricks.labs.ucx.hive_metastore.table_size import TableSizeCrawler
 from databricks.labs.ucx.workspace_access.generic import WorkspaceListing
 from databricks.labs.ucx.workspace_access.groups import GroupManager
 from databricks.labs.ucx.workspace_access.manager import PermissionManager
@@ -55,6 +56,17 @@ def crawl_grants(cfg: WorkspaceConfig):
     grants.snapshot()
 
 
+@task("assessment", depends_on=[crawl_tables])
+def estimate_table_size_for_migration(cfg: WorkspaceConfig):
+    """Scans the previously created Delta table named `$inventory_database.tables` and locate tables that cannot be
+    "synced". These tables will have to be cloned in the migration process.
+    Assesses the size of these tables and create `$inventory_database.table_size` table to list these sizes.
+    The table size is a factor in deciding whether to clone these tables."""
+    backend = RuntimeBackend()
+    table_size = TableSizeCrawler(backend, cfg.inventory_database)
+    table_size.snapshot()
+
+
 @task("assessment")
 def crawl_mounts(cfg: WorkspaceConfig):
     """Defines the scope of the _mount points_ intended for migration into Unity Catalog. As these objects are not

tests/unit/framework/mocks.py

@@ -35,7 +35,7 @@ def fetch(self, sql) -> Iterator[Row]:
         if self._rows:
             for pattern in self._rows.keys():
                 r = re.compile(pattern)
-                if r.match(sql):
+                if r.search(sql):
                     logger.debug(f"Found match: {sql}")
                     rows.extend(self._rows[pattern])
         logger.debug(f"Returning rows: {rows}")

tests/unit/hive_metastore/test_table_size.py

@@ -0,0 +1,36 @@
+import sys
+
+from databricks.labs.ucx.hive_metastore.table_size import TableSize, TableSizeCrawler
+from tests.unit.framework.mocks import MockBackend
+
+
+class SparkSession:
+    pass
+
+
+def test_table_size_crawler(mocker):
+    errors = {}
+    rows = {
+        "table_size": [],
+        "hive_metastore.inventory_database.tables": [
+            ("hive_metastore", "db1", "table1", "MANAGED", "DELTA", "dbfs:/location/table", None),
+            ("hive_metastore", "db1", "table2", "MANAGED", "DELTA", "/dbfs/location/table", None),
+            ("hive_metastore", "db1", "table3", "MANAGED", "DELTA", "dbfs:/mnt/location/table", None),
+            ("hive_metastore", "db1", "table4", "MANAGED", "DELTA", "s3:/location/table", None),
+            ("hive_metastore", "db1", "table5", "MANAGED", "DELTA", "/dbfs/mnt/location/table", None),
+            ("hive_metastore", "db1", "table6", "MANAGED", "DELTA", "/dbfs/databricks-datasets/location/table", None),
+            ("hive_metastore", "db1", "table7", "MANAGED", "DELTA", "dbfs:/databricks-datasets/location/table", None),
+            ("hive_metastore", "db1", "table8", "MANAGED", "DELTA", "/databricks-datasets/location/table", None),
+            ("hive_metastore", "db1", "view", "VIEW", "DELTA", None, "SELECT * FROM TABLE"),
+        ],
+        "SHOW DATABASES": [("db1",)],
+    }
+    backend = MockBackend(fails_on_first=errors, rows=rows)
+    pyspark_sql_session = mocker.Mock()
+    sys.modules["pyspark.sql.session"] = pyspark_sql_session
+    tsc = TableSizeCrawler(backend, "inventory_database")
+    tsc._spark._jsparkSession.table().queryExecution().analyzed().stats().sizeInBytes.side_effect = [100, 200, 300]
+    results = tsc.snapshot()
+    assert len(results) == 2
+    assert TableSize("hive_metastore", "db1", "table1", 100) in results
+    assert TableSize("hive_metastore", "db1", "table2", 200) in results

tests/unit/hive_metastore/test_tables.py

@@ -107,3 +107,18 @@ def test_tables_returning_error_when_describing():
     tc = TablesCrawler(backend, "default")
     results = tc._crawl()
     assert len(results) == 1
+
+
+def test_is_dbfs_root():
+    assert Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/somelocation/tablename").is_dbfs_root()
+    assert Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/somelocation/tablename").is_dbfs_root()
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="dbfs:/mnt/somelocation/tablename").is_dbfs_root()
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="/dbfs/mnt/somelocation/tablename").is_dbfs_root()
+    assert not Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="dbfs:/databricks-datasets/somelocation/tablename"
+    ).is_dbfs_root()
+    assert not Table(
+        "a", "b", "c", "MANAGED", "DELTA", location="/dbfs/databricks-datasets/somelocation/tablename"
+    ).is_dbfs_root()
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="s3:/somelocation/tablename").is_dbfs_root()
+    assert not Table("a", "b", "c", "MANAGED", "DELTA", location="adls:/somelocation/tablename").is_dbfs_root()