Repository: wandb/weave
Pull request #4130: feat(weave): Support emitting storage size in object details view (trace_server)
Author: login: chance-wnb
Created: 2025-04-14T22:45:43Z | Merged: 2025-04-21T22:03:50Z

## Description

<!--
Include reference to internal ticket "Fixes WB-NNNNN" and/or GitHub issue "Fixes #NNNN" (if applicable)
-->

## Summary
This PR adds support for displaying storage size information in the object details view and table statistics. It consists of mostly the backend changes, and minimum frontend changes to fit the interface changes related to object storage size only.

TLDR, I changed the following two APIs to include storage size information

```
/objs/query  ---> this will be used to retrieve each individual object version's storage size,
/table/query_stats_batch ---> this will be used to retrieve each table's storage size from all of its rows.
```

After this PR is merged. The internal PR https://github.com/wandb/core/pull/28783 should merge soon.

As a reference, the Frontend change that is going to utilize the API endpoint in this PR can be reviewed [here](https://github.com/wandb/weave/pull/4142).

## Changes
### Backend Changes
1. Enhanced object querying:
   - Added `include_storage_size` option to `ObjQueryReq`
   - Updated `ObjectMetadataQueryBuilder` to handle storage size inclusion
   - Modified query result formatting to include storage size when requested
   -
2. In order to prevent breaking existing clients (FE+SDK) due to tempering the existing endpoint of `table_stats_query`:
   - Created a new endpoint `table_stats_query_batch` which takes plural form parameter `digests`
   - Added `include_storage_size` flag to control storage size retrieval
   - Created new `make_table_stats_query_with_storage_size` function for efficient storage size queries


### Frontend Changes
   - Updated TypeScript interfaces to include new storage size fields for object queries

### Testing
1. Added new test cases:
   - `test_obj_query_with_storage_size_clickhouse` to verify object storage size functionality
   - `test_table_query_stats_with_storage_size` to verify table storage size retrieval
   - Updated existing tests in `test_table_query.py` suite to test against the `table_stats_query_batch` api.

2. Added SQL query tests:
   - Fixed breakage in the `test_objects_query_builder.py` with storage size related tests
   - Added SQL formatting assertions for better query validation


## Future Improvements
- Hook up the frontend.
- Remove the deprecated `table/query_stats` endpoint, and use the batch API `table/query_stats_batch` (Far future when the conditions meet)




Diff:
diff --git a/tests/trace/test_client_trace.py b/tests/trace/test_client_trace.py
index 4835fe82a03d..1def6f27c4a9 100644
--- a/tests/trace/test_client_trace.py
+++ b/tests/trace/test_client_trace.py
@@ -30,6 +30,7 @@
     get_weave_client,
     set_weave_client_global,
 )
+from weave.trace.refs import parse_uri
 from weave.trace.vals import MissingSelfInstanceError
 from weave.trace.weave_client import sanitize_object_name
 from weave.trace_server import trace_server_interface as tsi
@@ -3693,6 +3694,58 @@ def nest(input_ref: str):
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
 
 
+def test_obj_query_with_storage_size_clickhouse(client):
+    """Test querying objects with storage size information"""
+    if client_is_sqlite(client):
+        pytest.skip("Skipping test for sqlite clients")
+
+    # Create a test object with some data to ensure it has size
+    dataset = weave.Dataset(name="test_dataset", rows=[{"key": "value" * 1000}])
+    weave.publish(dataset)
+
+    # Query the object with storage size included
+    res = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=get_client_project_id(client),
+            include_storage_size=True,
+            filter={"object_ids": ["test_dataset"]},
+        )
+    )
+
+    assert len(res.objs) == 1
+    queried_obj = res.objs[0]
+
+    # Verify that storage size is present
+    assert queried_obj.size_bytes is not None
+    assert queried_obj.size_bytes == 270  # Should have some size due to the test data
+
+    # Test that a table is created and its size is correct
+    table_ref = parse_uri(queried_obj.val["rows"])
+    res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
+            project_id=client._project_id(),
+            digests=[table_ref.digest],
+            include_storage_size=True,
+        )
+    )
+
+    assert res.tables[0].storage_size_bytes == 5011
+
+    # Query without storage size (default behavior)
+    res_without_size = client.server.objs_query(
+        tsi.ObjQueryReq(
+            project_id=get_client_project_id(client),
+            filter={"object_ids": ["test_dataset"]},
+        )
+    )
+
+    assert len(res_without_size.objs) == 1
+    queried_obj_without_size = res_without_size.objs[0]
+
+    # Verify that storage size is not included when not requested
+    assert queried_obj_without_size.size_bytes is None
+
+
 def test_call_query_stream_with_costs_and_storage_size(client, clickhouse_client):
     if client_is_sqlite(client):
         # dont run this test for sqlite
diff --git a/tests/trace/test_objects_query_builder.py b/tests/trace/test_objects_query_builder.py
index 73fa99ecc527..b4f2b544e224 100644
--- a/tests/trace/test_objects_query_builder.py
+++ b/tests/trace/test_objects_query_builder.py
@@ -1,4 +1,5 @@
 import pytest
+import sqlparse
 
 from weave.trace_server import trace_server_interface as tsi
 from weave.trace_server.objects_query_builder import (
@@ -211,6 +212,13 @@ def test_object_query_builder_sort():
         FROM object_versions"""
 
 
+def assert_sql(exp_query, actual_query):
+    exp_formatted = sqlparse.format(exp_query, reindent=True)
+    actual_formatted = sqlparse.format(actual_query, reindent=True)
+
+    assert exp_formatted == actual_formatted
+
+
 def test_object_query_builder_metadata_query_basic():
     builder = ObjectMetadataQueryBuilder(project_id="test_project")
     builder.add_digests_conditions("latest")
@@ -222,11 +230,11 @@ def test_object_query_builder_metadata_query_basic():
         WHERE project_id = {{project_id: String}}
     )
     WHERE rn = 1
-)
+) as main
 WHERE ((is_latest = 1) AND (deleted_at IS NULL))
 ORDER BY created_at ASC"""
 
-    assert query == expected_query
+    assert_sql(query, expected_query)
     assert parameters == {"project_id": "test_project"}
 
 
@@ -250,13 +258,13 @@ def test_object_query_builder_metadata_query_with_limit_offset_sort():
         WHERE project_id = {{project_id: String}} AND object_id = {{object_id: String}}
     )
     WHERE rn = 1
-)
+) as main
 WHERE ((((digest = {{version_digest_0: String}}) OR (digest = {{version_digest_1: String}}) OR (version_index = {{version_index_2: Int64}}))) AND (base_object_class IN {{base_object_classes: Array(String)}}) AND (deleted_at IS NULL))
 ORDER BY created_at DESC
 LIMIT 10
 OFFSET 5"""
 
-    assert query == expected_query
+    assert_sql(query, expected_query)
     assert parameters == {
         "project_id": "test_project",
         "object_id": "object_1",
@@ -280,11 +288,11 @@ def test_objects_query_metadata_op():
         WHERE project_id = {{project_id: String}} AND object_id = {{object_id: String}}
     )
     WHERE rn = 1
-)
+) as main
 WHERE ((is_op = 1) AND (version_index = {{version_index_0: Int64}}) AND (deleted_at IS NULL))
 ORDER BY created_at ASC"""
 
-    assert query == expected_query
+    assert_sql(query, expected_query)
     assert parameters == {
         "project_id": "test_project",
         "object_id": "my_op",
@@ -310,7 +318,7 @@ def test_make_objects_val_query_and_parameters():
         GROUP BY object_id, digest
     """
 
-    assert query == expected_query
+    assert_sql(query, expected_query)
     assert parameters == {
         "project_id": "test_project",
         "object_ids": ["object_1"],
diff --git a/tests/trace/test_table_query.py b/tests/trace/test_table_query.py
index 1a703d6342b7..8dafc07a78be 100644
--- a/tests/trace/test_table_query.py
+++ b/tests/trace/test_table_query.py
@@ -1,6 +1,11 @@
 import random
 from collections.abc import Iterator
 
+import pytest
+
+from tests.trace.util import (
+    client_is_sqlite,
+)
 from weave.trace.weave_client import WeaveClient
 from weave.trace_server import trace_server_interface as tsi
 
@@ -265,40 +270,40 @@ def test_table_query_multiple_sort_criteria(client: WeaveClient):
 def test_table_query_stats(client: WeaveClient):
     digest, row_digests, data = generate_table_data(client, 10, 10)
 
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest=digest,
+            digests=[digest],
         )
     )
 
-    assert stats_res.count == len(data)
+    assert stats_res.tables[0].count == len(data)
 
 
 def test_table_query_stats_empty(client: WeaveClient):
     digest, row_digests, data = generate_table_data(client, 0, 0)
 
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest=digest,
+            digests=[digest],
         )
     )
 
-    assert stats_res.count == len(data)
+    assert stats_res.tables[0].count == len(data)
 
 
 def test_table_query_stats_missing(client: WeaveClient):
     digest, row_digests, data = generate_table_data(client, 10, 10)
 
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest="missing",
+            digests=["missing"],
         )
     )
 
-    assert stats_res.count == 0
+    assert len(stats_res.tables) == 0
 
 
 def generate_duplication_simple_table_data(
@@ -333,13 +338,13 @@ def test_table_query_with_duplicate_row_digests(client: WeaveClient):
             digest=res1["digest"],
         )
     )
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest=res1["digest"],
+            digests=[res1["digest"]],
         )
     )
-    assert len(res.rows) == stats_res.count == 10
+    assert len(res.rows) == stats_res.tables[0].count == 10
     assert [r.original_index for r in res.rows] == list(range(10))
 
     # Test filtered query for res1
@@ -360,13 +365,13 @@ def test_table_query_with_duplicate_row_digests(client: WeaveClient):
             digest=res2["digest"],
         )
     )
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest=res2["digest"],
+            digests=[res2["digest"]],
         )
     )
-    assert len(res.rows) == stats_res.count == 20
+    assert len(res.rows) == stats_res.tables[0].count == 20
     assert [r.original_index for r in res.rows] == list(range(20))
 
     # Test filtered query for res2
@@ -387,13 +392,13 @@ def test_table_query_with_duplicate_row_digests(client: WeaveClient):
             digest=res3["digest"],
         )
     )
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
             project_id=client._project_id(),
-            digest=res3["digest"],
+            digests=[res3["digest"]],
         )
     )
-    assert len(res.rows) == stats_res.count == 30
+    assert len(res.rows) == stats_res.tables[0].count == 30
     assert [r.original_index for r in res.rows] == list(range(30))
 
     # Test filtered query for res3
@@ -446,3 +451,21 @@ def test_duplicate_table_with_identical_rows(client: WeaveClient):
     # this is the same table, so we should get the same number of rows
     assert len(res.rows) == 10
     assert [r.original_index for r in res.rows] == list(range(10))
+
+
+def test_table_query_stats_with_storage_size(client: WeaveClient):
+    if client_is_sqlite(client):
+        pytest.skip("SQLite does not support storage size")
+
+    digest, row_digests, data = generate_table_data(client, 10, 10)
+
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
+            project_id=client._project_id(),
+            digests=[digest],
+            include_storage_size=True,
+        )
+    )
+
+    assert stats_res.tables[0].count == len(data)
+    assert stats_res.tables[0].storage_size_bytes > 0
diff --git a/weave/trace_server/clickhouse_schema.py b/weave/trace_server/clickhouse_schema.py
index d7e716f90034..f0112e3f665e 100644
--- a/weave/trace_server/clickhouse_schema.py
+++ b/weave/trace_server/clickhouse_schema.py
@@ -156,3 +156,4 @@ class SelectableCHObjSchema(BaseModel):
     version_index: int
     is_latest: int
     deleted_at: Optional[datetime.datetime] = None
+    size_bytes: Optional[int] = None
diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py
index 55dd8dee773a..347fdc9b5fa8 100644
--- a/weave/trace_server/clickhouse_trace_server_batched.py
+++ b/weave/trace_server/clickhouse_trace_server_batched.py
@@ -115,6 +115,7 @@
     VAL_DUMP_COLUMN_NAME,
     make_natural_sort_table_query,
     make_standard_table_query,
+    make_table_stats_query_with_storage_size,
 )
 from weave.trace_server.token_costs import (
     LLM_TOKEN_PRICES_TABLE,
@@ -745,8 +746,8 @@ def objs_query(self, req: tsi.ObjQueryReq) -> tsi.ObjQueryRes:
                 object_query_builder.add_order(sort.field, sort.direction)
         metadata_only = req.metadata_only or False
         object_query_builder.set_include_deleted(include_deleted=False)
+        object_query_builder.include_storage_size = req.include_storage_size or False
         objs = self._select_objs_query(object_query_builder, metadata_only)
-
         return tsi.ObjQueryRes(objs=[_ch_obj_to_obj_schema(obj) for obj in objs])
 
     def obj_delete(self, req: tsi.ObjDeleteReq) -> tsi.ObjDeleteRes:
@@ -1037,22 +1038,52 @@ def _table_query_stream(
                 digest=row[0], val=json.loads(row[1]), original_index=row[2]
             )
 
+    # This is a legacy endpoint, it should be removed once the client is mostly updated
     def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
+        batch_req = tsi.TableQueryStatsBatchReq(
+            project_id=req.project_id, digests=[req.digest]
+        )
+
+        res = self.table_query_stats_batch(batch_req)
+
+        if len(res.tables) != 1:
+            raise ValueError("Unexpected number of results", res)
+
+        count = res.tables[0].count
+        return tsi.TableQueryStatsRes(count=count)
+
+    def table_query_stats_batch(
+        self, req: tsi.TableQueryStatsBatchReq
+    ) -> tsi.TableQueryStatsBatchRes:
         parameters: dict[str, Any] = {
             "project_id": req.project_id,
-            "digest": req.digest,
+            "digests": req.digests,
         }
 
         query = """
-        SELECT length(row_digests)
+        SELECT digest, length(row_digests)
         FROM tables
-        WHERE project_id = {project_id:String} AND digest = {digest:String}
+        WHERE project_id = {project_id:String} AND digest IN {digests:Array(String)}
         """
 
+        if req.include_storage_size:
+            # Use an advanced query builder to get the storage size
+            pb = ParamBuilder()
+            query = make_table_stats_query_with_storage_size(
+                project_id=req.project_id,
+                table_digests=cast(list[str], req.digests),
+                pb=pb,
+            )
+            parameters = pb.get_params()
+
         query_result = self.ch_client.query(query, parameters=parameters)
-        count = query_result.result_rows[0][0] if query_result.result_rows else 0
 
-        return tsi.TableQueryStatsRes(count=count)
+        tables = [
+            _ch_table_stats_to_table_stats_schema(row)
+            for row in query_result.result_rows
+        ]
+
+        return tsi.TableQueryStatsBatchRes(tables=tables)
 
     def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
         # TODO: This reads one ref at a time, it should read them in batches
@@ -1951,7 +1982,9 @@ def _select_objs_query(
         obj_metadata_query = object_query_builder.make_metadata_query()
         parameters = object_query_builder.parameters or {}
         query_result = self._query_stream(obj_metadata_query, parameters)
-        metadata_result = format_metadata_objects_from_query_result(query_result)
+        metadata_result = format_metadata_objects_from_query_result(
+            query_result, object_query_builder.include_storage_size
+        )
 
         # -- Don't make second query for object values if metadata_only --
         if metadata_only or len(metadata_result) == 0:
@@ -2261,6 +2294,21 @@ def _ch_obj_to_obj_schema(ch_obj: SelectableCHObjSchema) -> tsi.ObjSchema:
         kind=ch_obj.kind,
         base_object_class=ch_obj.base_object_class,
         val=json.loads(ch_obj.val_dump),
+        size_bytes=ch_obj.size_bytes,
+    )
+
+
+def _ch_table_stats_to_table_stats_schema(
+    ch_table_stats_row: Sequence[Any],
+) -> tsi.TableStatsRow:
+    digest, count, storage_size_bytes = (lambda a, b, c=cast(Any, None): (a, b, c))(
+        *ch_table_stats_row
+    )
+
+    return tsi.TableStatsRow(
+        count=count,
+        digest=digest,
+        storage_size_bytes=storage_size_bytes,
     )
 
 
diff --git a/weave/trace_server/external_to_internal_trace_server_adapter.py b/weave/trace_server/external_to_internal_trace_server_adapter.py
index a70902292e10..88c8561ba683 100644
--- a/weave/trace_server/external_to_internal_trace_server_adapter.py
+++ b/weave/trace_server/external_to_internal_trace_server_adapter.py
@@ -289,10 +289,17 @@ def table_query_stream(
             self._internal_trace_server.table_query_stream, req
         )
 
+    # This is a legacy endpoint, it should be removed once the client is mostly updated
     def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
         req.project_id = self._idc.ext_to_int_project_id(req.project_id)
         return self._ref_apply(self._internal_trace_server.table_query_stats, req)
 
+    def table_query_stats_batch(
+        self, req: tsi.TableQueryStatsBatchReq
+    ) -> tsi.TableQueryStatsBatchRes:
+        req.project_id = self._idc.ext_to_int_project_id(req.project_id)
+        return self._ref_apply(self._internal_trace_server.table_query_stats_batch, req)
+
     def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
         return self._ref_apply(self._internal_trace_server.refs_read_batch, req)
 
diff --git a/weave/trace_server/objects_query_builder.py b/weave/trace_server/objects_query_builder.py
index 05cca2aae8aa..c7a4aa063183 100644
--- a/weave/trace_server/objects_query_builder.py
+++ b/weave/trace_server/objects_query_builder.py
@@ -85,13 +85,21 @@ def _make_object_id_conditions_part(
 
 def format_metadata_objects_from_query_result(
     query_result: Iterator[tuple[Any, ...]],
+    include_storage_size: bool = False,
 ) -> list[SelectableCHObjSchema]:
     result = []
     for row in query_result:
         # Add an empty val_dump to the end of the row
         row_with_val_dump = row + ("{}",)
-        columns_with_val_dump = OBJECT_METADATA_COLUMNS + ["val_dump"]
+        columns_with_val_dump = list(OBJECT_METADATA_COLUMNS)
+
+        if include_storage_size:
+            columns_with_val_dump += ["size_bytes"]
+
+        columns_with_val_dump += ["val_dump"]
+
         row_dict = dict(zip(columns_with_val_dump, row_with_val_dump))
+
         row_model = SelectableCHObjSchema.model_validate(row_dict)
         result.append(row_model)
     return result
@@ -116,6 +124,7 @@ def __init__(
         self._offset: Optional[int] = None
         self._sort_by: list[tsi.SortBy] = []
         self._include_deleted: bool = include_deleted
+        self.include_storage_size: bool = False
 
     @property
     def conditions_part(self) -> str:
@@ -242,10 +251,25 @@ def set_include_deleted(self, include_deleted: bool) -> None:
         self._include_deleted = include_deleted
 
     def make_metadata_query(self) -> str:
-        columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+        columns = list(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
+            columns += ["size_bytes"]
+        columns_str = ",\n    ".join(columns)
+
+        join_clause = ""
+        if self.include_storage_size:
+            join_clause = f"""
+            LEFT JOIN (
+                SELECT * FROM object_versions_stats WHERE object_versions_stats.project_id = {{project_id: String}}
+            ) as object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest
+            """
+
         query = f"""
 SELECT
-    {columns}
+    {columns_str}
 FROM (
     SELECT
         project_id,
@@ -308,7 +332,9 @@ def make_metadata_query(self) -> str:
         WHERE project_id = {{project_id: String}}{self.object_id_conditions_part}
     )
     WHERE rn = 1
-)"""
+) as {main_table_alias}
+    {join_clause}
+"""
         if self.conditions_part:
             query += f"\n{self.conditions_part}"
         if self.sort_part:
diff --git a/weave/trace_server/sqlite_trace_server.py b/weave/trace_server/sqlite_trace_server.py
index 0c5ea5e4423b..284a402925a7 100644
--- a/weave/trace_server/sqlite_trace_server.py
+++ b/weave/trace_server/sqlite_trace_server.py
@@ -1127,26 +1127,49 @@ def table_query(self, req: tsi.TableQueryReq) -> tsi.TableQueryRes:
             ]
         )
 
+    # This is a legacy endpoint, it should be removed once the client is mostly updated
     def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
-        parameters: list[Any] = [req.project_id, req.digest]
+        batch_req = tsi.TableQueryStatsBatchReq(
+            project_id=req.project_id, digests=[req.digest]
+        )
+
+        res = self.table_query_stats_batch(batch_req)
+
+        if len(res.tables) != 1:
+            raise ValueError("Unexpected number of results", res)
+
+        count = res.tables[0].count
+        return tsi.TableQueryStatsRes(count=count)
+
+    def table_query_stats_batch(
+        self, req: tsi.TableQueryStatsBatchReq
+    ) -> tsi.TableQueryStatsBatchRes:
+        parameters: list[Any] = [req.project_id] + list(req.digests or [])
+
+        placeholders = ",".join(["?" for _ in (req.digests or [])])
 
-        query = """
-        SELECT json_array_length(row_digests)
+        query = f"""
+        SELECT digest, json_array_length(row_digests)
         FROM
             tables
         WHERE
             tables.project_id = ? AND
-            tables.digest = ?
+            tables.digest in ({placeholders})
         """
 
         conn, cursor = get_conn_cursor(self.db_path)
         cursor.execute(query, parameters)
-        row = cursor.fetchone()
-        count = 0
-        if row is not None:
-            count = row[0]
 
-        return tsi.TableQueryStatsRes(count=count)
+        query_result = cursor.fetchall()
+
+        tables = []
+
+        for row in query_result:
+            count = row[1]
+            digest = row[0]
+            tables.append(tsi.TableStatsRow(count=count, digest=digest))
+
+        return tsi.TableQueryStatsBatchRes(tables=tables)
 
     def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
         # TODO: This reads one ref at a time, it should read them in batches
diff --git a/weave/trace_server/table_query_builder.py b/weave/trace_server/table_query_builder.py
index 038718eb80b0..ab3f5398c2f7 100644
--- a/weave/trace_server/table_query_builder.py
+++ b/weave/trace_server/table_query_builder.py
@@ -113,3 +113,29 @@ def make_standard_table_query(
     {sql_safe_offset}
     """
     return query
+
+
+def make_table_stats_query_with_storage_size(
+    project_id: str,
+    table_digests: list[str],
+    pb: ParamBuilder,
+) -> str:
+    """Generate a query for table stats with storage size and length(num of rows)."""
+    project_id_name = pb.add_param(project_id)
+    digest_ids = pb.add_param(table_digests)
+
+    query = f"""
+    SELECT tb_digest, any(length), sum(size_bytes) FROM
+    (
+        SELECT digest as tb_digest, length(row_digests) as length, row_digests
+        FROM tables
+        WHERE project_id = {{{project_id_name}: String}} AND digest in {{{digest_ids}: Array(String)}}
+    ) ARRAY JOIN row_digests as row_digest
+    LEFT JOIN
+    (
+        SELECT * FROM table_rows_stats WHERE table_rows_stats.project_id = {{{project_id_name}: String}}
+    ) as table_rows_stats ON table_rows_stats.digest = row_digest
+
+    GROUP BY tb_digest
+    """
+    return query
diff --git a/weave/trace_server/trace_server_interface.py b/weave/trace_server/trace_server_interface.py
index 125a3b4bd3d5..cad1949257e8 100644
--- a/weave/trace_server/trace_server_interface.py
+++ b/weave/trace_server/trace_server_interface.py
@@ -198,6 +198,7 @@ class ObjSchema(BaseModel):
     val: Any
 
     wb_user_id: Optional[str] = Field(None, description=WB_USER_ID_DESCRIPTION)
+    size_bytes: Optional[int] = None
 
 
 class ObjSchemaForInsert(BaseModel):
@@ -547,6 +548,10 @@ class ObjQueryReq(BaseModel):
         description="If true, the `val` column is not read from the database and is empty."
         "All other fields are returned.",
     )
+    include_storage_size: Optional[bool] = Field(
+        default=False,
+        description="If true, the `size_bytes` column is returned.",
+    )
 
 
 class ObjDeleteReq(BaseModel):
@@ -751,7 +756,25 @@ class TableQueryStatsReq(BaseModel):
     )
     digest: str = Field(
         description="The digest of the table to query",
-        examples=["aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims"],
+    )
+
+
+class TableQueryStatsBatchReq(BaseModel):
+    project_id: str = Field(
+        description="The ID of the project", examples=["my_entity/my_project"]
+    )
+
+    digests: Optional[list[str]] = Field(
+        description="The digests of the tables to query",
+        examples=[
+            "aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims",
+            "smirva431etnsroatsratlrampgrmeangmpr5344aplatmipa31ltvsmiераnoa",
+        ],
+        default=[],
+    )
+    include_storage_size: Optional[bool] = Field(
+        default=False,
+        description="If true, the `storage_size_bytes` column is returned.",
     )
 
 
@@ -759,6 +782,16 @@ class TableQueryStatsRes(BaseModel):
     count: int
 
 
+class TableStatsRow(BaseModel):
+    count: int
+    digest: str
+    storage_size_bytes: Optional[int] = None
+
+
+class TableQueryStatsBatchRes(BaseModel):
+    tables: list[TableStatsRow]
+
+
 class RefsReadBatchReq(BaseModel):
     refs: list[str]
 
@@ -1005,6 +1038,9 @@ def table_update(self, req: TableUpdateReq) -> TableUpdateRes: ...
     def table_query(self, req: TableQueryReq) -> TableQueryRes: ...
     def table_query_stream(self, req: TableQueryReq) -> Iterator[TableRowSchema]: ...
     def table_query_stats(self, req: TableQueryStatsReq) -> TableQueryStatsRes: ...
+    def table_query_stats_batch(
+        self, req: TableQueryStatsBatchReq
+    ) -> TableQueryStatsBatchRes: ...
 
     # Ref API
     def refs_read_batch(self, req: RefsReadBatchReq) -> RefsReadBatchRes: ...
diff --git a/weave/trace_server_bindings/caching_middleware_trace_server.py b/weave/trace_server_bindings/caching_middleware_trace_server.py
index 7931fbdbdc4c..425a237ab669 100644
--- a/weave/trace_server_bindings/caching_middleware_trace_server.py
+++ b/weave/trace_server_bindings/caching_middleware_trace_server.py
@@ -280,6 +280,7 @@ def table_query_stream(
         # I am not sure the best way to cache the iterator here. TODO
         return self._next_trace_server.table_query_stream(req)
 
+    # This is a legacy endpoint, it should be removed once the client is mostly updated
     def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsRes:
         if not digest_is_cacheable(req.digest):
             return self._next_trace_server.table_query_stats(req)
@@ -287,6 +288,17 @@ def table_query_stats(self, req: tsi.TableQueryStatsReq) -> tsi.TableQueryStatsR
             self._next_trace_server.table_query_stats, req, tsi.TableQueryStatsRes
         )
 
+    def table_query_stats_batch(
+        self, req: tsi.TableQueryStatsBatchReq
+    ) -> tsi.TableQueryStatsBatchRes:
+        if any(not digest_is_cacheable(digest) for digest in req.digests or []):
+            return self._next_trace_server.table_query_stats_batch(req)
+        return self._with_cache_pydantic(
+            self._next_trace_server.table_query_stats_batch,
+            req,
+            tsi.TableQueryStatsBatchRes,
+        )
+
     def refs_read_batch(self, req: tsi.RefsReadBatchReq) -> tsi.RefsReadBatchRes:
         """Read multiple refs, utilizing cache for individual refs.
 
diff --git a/weave/trace_server_bindings/remote_http_trace_server.py b/weave/trace_server_bindings/remote_http_trace_server.py
index 23457656a133..956df290bc2e 100644
--- a/weave/trace_server_bindings/remote_http_trace_server.py
+++ b/weave/trace_server_bindings/remote_http_trace_server.py
@@ -451,6 +451,16 @@ def table_query_stats(
             "/table/query_stats", req, tsi.TableQueryStatsReq, tsi.TableQueryStatsRes
         )
 
+    def table_query_stats_batch(
+        self, req: Union[tsi.TableQueryStatsReq, dict[str, Any]]
+    ) -> tsi.TableQueryStatsRes:
+        return self._generic_request(
+            "/table/query_stats_batch",
+            req,
+            tsi.TableQueryStatsBatchReq,
+            tsi.TableQueryStatsBatchRes,
+        )
+
     def refs_read_batch(
         self, req: Union[tsi.RefsReadBatchReq, dict[str, Any]]
     ) -> tsi.RefsReadBatchRes:


Reviews:
user: login: neutralino1
body: Thank you.
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-15T16:10:53Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-15T17:40:16Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T23:40:13Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T23:50:24Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:17:50Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:27:05Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:29:47Z
user: login: tssweeney
body: Very lovely implementation! Nice job (: 

Tactically, we probably want:
1. A quick eye from @gtarpenning regarding the query builder
2. Split the python and typescript since the typescript can't be landed before the server picks up the changes
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-17T03:34:20Z
user: login: gtarpenning
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-17T16:04:37Z
user: login: gtarpenning
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-17T16:09:07Z
user: login: gtarpenning
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-17T16:10:07Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T22:07:38Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T22:37:06Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T07:47:23Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-18T18:48:44Z
user: login: gtarpenning
body: lgtm, thanks for the query change. Let's make sure to test in a QA environment before deploying just to be sure that the shape of real data doesn't cause any issues. I don't expect it to, but this is a sufficiently large change as to make me slightly cautious. 
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-18T19:27:19Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-19T03:38:43Z
user: login: tssweeney
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-19T03:39:22Z
user: login: tssweeney
body: Approving with some lightweight comments on additional tests
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-19T03:40:40Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T19:22:30Z
user: login: chance-wnb
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-19T19:22:50Z

Review comments:
diff_hunk: @@ -749,14 +754,33 @@ class TableQueryStatsReq(BaseModel):
     project_id: str = Field(
         description="The ID of the project", examples=["my_entity/my_project"]
     )
-    digest: str = Field(
-        description="The digest of the table to query",
+    digest: Optional[str] = Field(
+        description="The old singular digest of the table to query, this is temporary and will be removed in a future release",
         examples=["aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims"],
+        default=None,
+    )
+    digests: Optional[list[str]] = Field(
+        description="The digests of the tables to query",
+        examples=[
+            "aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims",
+            "smirva431etnsroatsratlrampgrmeangmpr5344aplatmipa31ltvsmiераnoa",
+        ],
+        default=[],
+    )
+    include_storage_size: Optional[bool] = Field(
+        default=False,
+        description="If true, the `storage_size_bytes` column is returned.",
     )
 
 
-class TableQueryStatsRes(BaseModel):
+class TableStatsResRow(BaseModel):
     count: int
+    digest: str
+    storage_size_bytes: Optional[int] = None
+
+
+class TableQueryStatsRes(BaseModel):
path: weave/trace_server/trace_server_interface.py
user: login: tssweeney
body: this introduces a breaking change to the existing service layer. Specifically: this will change the http spec. there are two problems here:
1. practically, from a deployment perspective, the gap between deploying the new frontend and backend will have a broken UI
2. any users of this stats endpoint directly will break. There is low likelihood that end users are directly targeting this, so it might be a risk we are willing to take, but certainly not best practice.

Either: a) we should make the changes such that the are not breaking; or b) introduce a batch version of the endpoint
created_at: 2025-04-15T17:40:16Z
author_association: COLLABORATOR
diff_hunk: @@ -749,14 +754,33 @@ class TableQueryStatsReq(BaseModel):
     project_id: str = Field(
         description="The ID of the project", examples=["my_entity/my_project"]
     )
-    digest: str = Field(
-        description="The digest of the table to query",
+    digest: Optional[str] = Field(
+        description="The old singular digest of the table to query, this is temporary and will be removed in a future release",
         examples=["aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims"],
+        default=None,
+    )
+    digests: Optional[list[str]] = Field(
+        description="The digests of the tables to query",
+        examples=[
+            "aonareimsvtl13apimtalpa4435rpmgnaemrpgmarltarstaorsnte134avrims",
+            "smirva431etnsroatsratlrampgrmeangmpr5344aplatmipa31ltvsmiераnoa",
+        ],
+        default=[],
+    )
+    include_storage_size: Optional[bool] = Field(
+        default=False,
+        description="If true, the `storage_size_bytes` column is returned.",
     )
 
 
-class TableQueryStatsRes(BaseModel):
+class TableStatsResRow(BaseModel):
     count: int
+    digest: str
+    storage_size_bytes: Optional[int] = None
+
+
+class TableQueryStatsRes(BaseModel):
path: weave/trace_server/trace_server_interface.py
user: login: chance-wnb
body: Yes, all agreed.

On
> any users of this stats endpoint directly will break. There is low likelihood that end users are directly targeting this, so it might be a risk we are willing to take, but certainly not best practice.

Just FYI, this API endpoint is used in the dataset read operation, which actually has a high impact.
created_at: 2025-04-15T23:40:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -265,40 +270,40 @@ def test_table_query_multiple_sort_criteria(client: WeaveClient):
 def test_table_query_stats(client: WeaveClient):
     digest, row_digests, data = generate_table_data(client, 10, 10)
 
-    stats_res = client.server.table_query_stats(
-        tsi.TableQueryStatsReq(
+    stats_res = client.server.table_query_stats_batch(
path: tests/trace/test_table_query.py
user: login: chance-wnb
body: So I am changing this suite to test `TableQueryStatsBatchReq` instead of `TableQueryStatsReq` and regard `TableQueryStatsReq` as obsolete.
created_at: 2025-04-15T23:50:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -3691,3 +3691,45 @@ def nest(input_ref: str):
     assert calls[1].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].output["woahhhh"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
+
+
+def test_obj_query_with_storage_size_clickhouse(client):
+    """Test querying objects with storage size information"""
+    if client_is_sqlite(client):
+        pytest.skip("Skipping test for sqlite clients")
+
+    # Create a test object with some data to ensure it has size
+    dataset = weave.Dataset(name="test_dataset", rows=[{"key": "value" * 1000}])
+    weave.publish(dataset)
+
+    print(dataset)
path: tests/trace/test_client_trace.py
user: login: tssweeney
body: dev tracks
created_at: 2025-04-17T03:17:50Z
author_association: COLLABORATOR
diff_hunk: @@ -113,3 +113,25 @@ def make_standard_table_query(
     {sql_safe_offset}
     """
     return query
+
+
+def make_table_stats_query_with_storage_size(
+    project_id: str,
+    table_digests: list[str],
+    pb: ParamBuilder,
+) -> str:
+    """Generate a query for table stats with storage size and length(num of rows)."""
+    project_id_name = pb.add_param(project_id)
+    digest_ids = pb.add_param(table_digests)
+
+    query = f"""
+    SELECT tb_digest, any(length), sum(size_bytes) FROM
+    (
+        SELECT digest as tb_digest, length(row_digests) as length, row_digests
+        FROM tables
+        WHERE project_id = {{{project_id_name}: String}} AND digest in {{{digest_ids}: Array(String)}}
+    ) ARRAY JOIN row_digests as row_digest
+    LEFT JOIN table_rows_stats on table_rows_stats.digest = row_digest
path: weave/trace_server/table_query_builder.py
user: login: tssweeney
body: it might improve perf to have a project_id in this part of the left join as well
created_at: 2025-04-17T03:27:05Z
author_association: COLLABORATOR
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
+            columns += ",\n   size_bytes"
+
+        join_clause = ""
+        if self.include_storage_size:
+            join_clause = f"LEFT JOIN object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest"
path: weave/trace_server/objects_query_builder.py
user: login: tssweeney
body: might want to include project id as well
created_at: 2025-04-17T03:29:47Z
author_association: COLLABORATOR
diff_hunk: @@ -3691,3 +3691,45 @@ def nest(input_ref: str):
     assert calls[1].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].output["woahhhh"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
+
+
+def test_obj_query_with_storage_size_clickhouse(client):
+    """Test querying objects with storage size information"""
+    if client_is_sqlite(client):
+        pytest.skip("Skipping test for sqlite clients")
+
+    # Create a test object with some data to ensure it has size
+    dataset = weave.Dataset(name="test_dataset", rows=[{"key": "value" * 1000}])
+    weave.publish(dataset)
+
+    print(dataset)
path: tests/trace/test_client_trace.py
user: login: gtarpenning
body: the prints are only collected on test failure, I sort of like them in there. If the test fails you get more info, if it passes no harm. 
created_at: 2025-04-17T16:04:37Z
author_association: MEMBER
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
path: weave/trace_server/objects_query_builder.py
user: login: gtarpenning
body: nit: could add "size_bytes" to a `columns_list` var and only combine to str once. like: 
```python
columns = OBJECT_METADATA_COLUMNS
if self.include_storage_size:
    columns += ["size_bytes"]
column_str = ",\n ".join(columns)
```
created_at: 2025-04-17T16:09:07Z
author_association: MEMBER
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
+            columns += ",\n   size_bytes"
+
+        join_clause = ""
+        if self.include_storage_size:
+            join_clause = f"LEFT JOIN object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest"
path: weave/trace_server/objects_query_builder.py
user: login: gtarpenning
body: we can run an `EXPLAIN indexes=1` to see if that's required, but yeah doesn't hurt. 
created_at: 2025-04-17T16:10:06Z
author_association: MEMBER
diff_hunk: @@ -3691,3 +3691,45 @@ def nest(input_ref: str):
     assert calls[1].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].output["woahhhh"]["hi"]["there"]["foo"] == "bar"
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
+
+
+def test_obj_query_with_storage_size_clickhouse(client):
+    """Test querying objects with storage size information"""
+    if client_is_sqlite(client):
+        pytest.skip("Skipping test for sqlite clients")
+
+    # Create a test object with some data to ensure it has size
+    dataset = weave.Dataset(name="test_dataset", rows=[{"key": "value" * 1000}])
+    weave.publish(dataset)
+
+    print(dataset)
path: tests/trace/test_client_trace.py
user: login: chance-wnb
body: Yeah, thanks, my bad, and I don't think the logging for the created dataset is that useful. I will just remove.
created_at: 2025-04-17T22:07:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
path: weave/trace_server/objects_query_builder.py
user: login: chance-wnb
body: Yeah, much better. Made it that way, thanks!
created_at: 2025-04-17T22:37:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
+            columns += ",\n   size_bytes"
+
+        join_clause = ""
+        if self.include_storage_size:
+            join_clause = f"LEFT JOIN object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest"
path: weave/trace_server/objects_query_builder.py
user: login: chance-wnb
body: Thanks for the suggestion. I made it to
```sql
LEFT JOIN object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest
            AND object_versions_stats.project_id = {{project_id: String}}
```

however, maybe what you suggested is 

```sql
LEFT JOIN (
        SELECT * FROM object_versions_stats WHERE object_versions_stats.project_id =  {{project_id: String}}
    ) as filtered_object_versions_stats ON filtered_object_versions_stats.digest = main.digest
```

let me know if the latter is preferred.
created_at: 2025-04-18T07:47:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -243,6 +252,16 @@ def set_include_deleted(self, include_deleted: bool) -> None:
 
     def make_metadata_query(self) -> str:
         columns = ",\n    ".join(OBJECT_METADATA_COLUMNS)
+
+        main_table_alias = "main"
+
+        if self.include_storage_size:
+            columns += ",\n   size_bytes"
+
+        join_clause = ""
+        if self.include_storage_size:
+            join_clause = f"LEFT JOIN object_versions_stats ON object_versions_stats.digest = {main_table_alias}.digest"
path: weave/trace_server/objects_query_builder.py
user: login: chance-wnb
body: I adopted the 2nd approach(sub-query) now.
created_at: 2025-04-18T18:48:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -3693,6 +3693,46 @@ def nest(input_ref: str):
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
 
 
+def test_obj_query_with_storage_size_clickhouse(client):
path: tests/trace/test_client_trace.py
user: login: tssweeney
body: this test seems like a prime opportunity to also assert on the table query size
created_at: 2025-04-19T03:38:43Z
author_association: COLLABORATOR
diff_hunk: @@ -446,3 +451,21 @@ def test_duplicate_table_with_identical_rows(client: WeaveClient):
     # this is the same table, so we should get the same number of rows
     assert len(res.rows) == 10
     assert [r.original_index for r in res.rows] == list(range(10))
+
+
+def test_table_query_stats_with_storage_size(client: WeaveClient):
+    if client_is_sqlite(client):
+        pytest.skip("SQLite does not support storage size")
+
+    digest, row_digests, data = generate_table_data(client, 10, 10)
+
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
+            project_id=client._project_id(),
+            digests=[digest],
+            include_storage_size=True,
+        )
+    )
+
+    assert stats_res.tables[0].count == len(data)
+    assert stats_res.tables[0].storage_size_bytes > 0
path: tests/trace/test_table_query.py
user: login: tssweeney
body: why not have this exact?
created_at: 2025-04-19T03:39:22Z
author_association: COLLABORATOR
diff_hunk: @@ -446,3 +451,21 @@ def test_duplicate_table_with_identical_rows(client: WeaveClient):
     # this is the same table, so we should get the same number of rows
     assert len(res.rows) == 10
     assert [r.original_index for r in res.rows] == list(range(10))
+
+
+def test_table_query_stats_with_storage_size(client: WeaveClient):
+    if client_is_sqlite(client):
+        pytest.skip("SQLite does not support storage size")
+
+    digest, row_digests, data = generate_table_data(client, 10, 10)
+
+    stats_res = client.server.table_query_stats_batch(
+        tsi.TableQueryStatsBatchReq(
+            project_id=client._project_id(),
+            digests=[digest],
+            include_storage_size=True,
+        )
+    )
+
+    assert stats_res.tables[0].count == len(data)
+    assert stats_res.tables[0].storage_size_bytes > 0
path: tests/trace/test_table_query.py
user: login: chance-wnb
body: `generate_table_data(client, 10, 10)` generates random data. so there is no exact number to test against. `generate_table_data()` is pretty good for other test cases, so I decide to not temper it. 

Considering that in the `test_client_trace` unit test I also added logic to verify the table size, and that test scenario already tests the exact number. I am not over-complicating things here. 
created_at: 2025-04-19T19:22:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -3693,6 +3693,46 @@ def nest(input_ref: str):
     assert calls[2].inputs["input_ref"]["hi"]["there"]["foo"] == "bar"
 
 
+def test_obj_query_with_storage_size_clickhouse(client):
path: tests/trace/test_client_trace.py
user: login: chance-wnb
body: Yep, added that. Thanks for noticing the opportunity!
created_at: 2025-04-19T19:22:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: circle-job-mirror[bot]
created_at: 2025-04-14T23:16:43Z
body: Preview this PR with FeatureBee: <https://beta.wandb.ai/?betaVersion=b90b9822d69d4e0191181a6cd1911ae63f5b61b7>
author_association: NONE
user: login: chance-wnb
created_at: 2025-04-18T07:47:47Z
body: > Very lovely implementation! Nice job (:
> 
> Tactically, we probably want:
> 
> 1. A quick eye from @gtarpenning regarding the query builder
> 2. Split the python and typescript since the typescript can't be landed before the server picks up the changes

Thank you so much! I have removed the TS changes from this PR.
author_association: CONTRIBUTOR
user: login: tssweeney
created_at: 2025-04-19T03:41:25Z
body: Also, i would normally like to break this into 2 prs since there are 2 distinct functionalities produced and tested. Ok to leave as is, but in the future consider breaking up more.
author_association: COLLABORATOR
user: login: chance-wnb
created_at: 2025-04-19T18:42:57Z
body: > Also, i would normally like to break this into 2 prs since there are 2 distinct functionalities produced and tested. Ok to leave as is, but in the future consider breaking up more.

That's a completely valid point and I am actually aware of that. Due to leadership's urgent expectations for wrapping up the project, and potential back-n-forth during code reviews if splitting, I had to resort to this desperate measure. 

I will absolutely try to minimize the occurrences like this in the future.
author_association: CONTRIBUTOR
user: login: chance-wnb
created_at: 2025-04-21T21:33:09Z
body: > lgtm, thanks for the query change. Let's make sure to test in a QA environment before deploying just to be sure that the shape of real data doesn't cause any issues. I don't expect it to, but this is a sufficiently large change as to make me slightly cautious.

I ran the requests against the QA "weave-team/types-chance" project, the current scenarios all work fine. I also simulated requests as if they need to request storage size (since currently UI is not hooked up), the requests all return quickly. 
author_association: CONTRIBUTOR
user: login: gtarpenning
created_at: 2025-04-21T22:16:28Z
body: 🚢🚢🚢 
author_association: MEMBER