Add get_schema_at to SchemaHelperMixin and use it in DataSliceManagerView.

As a result, view.get_schema() does not have to re-analyze the whole schema of the manager anymore.

PiperOrigin-RevId: 859109377
Change-Id: I740ca79477b29bc06bbb10c33159fa17a42bf7bf

Author: svsgoogle

py/koladata/ext/persisted_data/bare_root_initial_data_manager_test.py

@@ -333,6 +333,25 @@ def test_copy(self):
     self.assertIs(copy._root_item, manager._root_item)
     self.assertIs(copy._schema_helper, manager._schema_helper)
 
+  def test_get_schema_at(self):
+    root_schema = kd.named_schema('root_schema')
+    manager = BareRootInitialDataManager.create_new(root_schema.new())
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(
+            data_slice_path_lib.DataSlicePath.parse_from_string('')
+        ),
+        root_schema,
+        ids_equality=True,
+    )
+
+    manager = BareRootInitialDataManager.create_new()
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(
+            data_slice_path_lib.DataSlicePath.parse_from_string('')
+        ),
+        kd.schema.new_schema(),
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/ext/persisted_data/data_slice_manager_interface.py

@@ -27,6 +27,12 @@ def get_schema(self) -> kd.types.DataSlice:
     """Returns the schema of the entire DataSlice managed by this manager."""
     raise NotImplementedError(type(self))
 
+  def get_schema_at(
+      self, path: data_slice_path_lib.DataSlicePath
+  ) -> kd.types.SchemaItem:
+    """Returns the schema of the DataSlice at the given path."""
+    raise NotImplementedError(type(self))
+
   def generate_paths(
       self, *, max_depth: int
   ) -> Generator[data_slice_path_lib.DataSlicePath, None, None]:

py/koladata/ext/persisted_data/data_slice_manager_view.py

@@ -93,14 +93,7 @@ def get_schema(self) -> kd.types.SchemaItem:
     The view path must be valid, i.e. self.is_view_valid() must be True.
     """
     self._check_path_from_root_is_valid()
-    schema_helper = schema_helper_lib.SchemaHelper(
-        self._data_slice_manager.get_schema()
-    )
-    return schema_helper.get_subschema_at(
-        schema_helper.get_schema_node_name_for_data_slice_path(
-            self._path_from_root
-        )
-    )
+    return self._data_slice_manager.get_schema_at(self._path_from_root)
 
   def get_data_slice(
       self, *, with_ancestors: bool = False, with_descendants: bool = False

py/koladata/ext/persisted_data/initial_data_manager_interface.py

@@ -69,6 +69,12 @@ def get_schema(self) -> kd.types.SchemaItem:
     """
     raise NotImplementedError(type(self))
 
+  def get_schema_at(
+      self, path: data_slice_path_lib.DataSlicePath
+  ) -> kd.types.SchemaItem:
+    """Returns the schema of the DataSlice at the given path."""
+    raise NotImplementedError(type(self))
+
   def get_all_schema_node_names(self) -> AbstractSet[str]:
     """Returns all the schema node names of self.get_schema().
 

py/koladata/ext/persisted_data/persisted_incremental_data_slice_manager_test.py

@@ -4620,6 +4620,57 @@ def test_description_of_initial_revision(self):
         ['Some bespoke description'],
     )
 
+  def test_get_schema_at(self):
+    persistence_dir = os.path.join(
+        self.create_tempdir().full_path, 'persisted_dataslice'
+    )
+    manager = PersistedIncrementalDataSliceManager.create_new(persistence_dir)
+    query_schema = kd.named_schema('query')
+    manager.update(
+        at_path=parse_dsp(''),
+        attr_name='query',
+        attr_value=kd.list([
+            query_schema.new(query_id='q1'),
+            query_schema.new(query_id='q2'),
+        ]),
+    )
+    doc_schema = kd.named_schema('doc')
+    manager.update(
+        at_path=parse_dsp('.query[:]'),
+        attr_name='doc',
+        attr_value=kd.slice([
+            doc_schema.new(doc_id=kd.slice([0, 1, 2, 3])).implode(),
+            doc_schema.new(doc_id=kd.slice([4, 5, 6])).implode(),
+        ]),
+    )
+
+    full_schema = manager.get_schema()
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(parse_dsp('')),
+        full_schema,
+        ids_equality=True,
+    )
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(parse_dsp('.query[:]')),
+        full_schema.query.get_item_schema(),
+        ids_equality=True,
+    )
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(parse_dsp('.query[:].query_id')),
+        kd.STRING,
+        ids_equality=True,
+    )
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(parse_dsp('.query[:].doc[:]')),
+        full_schema.query.get_item_schema().doc.get_item_schema(),
+        ids_equality=True,
+    )
+    kd.testing.assert_equivalent(
+        manager.get_schema_at(parse_dsp('.query[:].doc[:].doc_id')),
+        kd.INT32,
+        ids_equality=True,
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

py/koladata/ext/persisted_data/schema_helper_mixin.py

@@ -18,7 +18,7 @@
 
 from koladata import kd
 from koladata.ext.persisted_data import data_slice_path as data_slice_path_lib
-from koladata.ext.persisted_data import schema_helper
+from koladata.ext.persisted_data import schema_helper as schema_helper_lib
 
 
 class SchemaHelperMixin:
@@ -31,7 +31,7 @@ class SchemaHelperMixin:
   get_data_slice_at(), exists() and get_schema() methods.
   """
 
-  def _get_schema_helper(self) -> schema_helper.SchemaHelper:
+  def _get_schema_helper(self) -> schema_helper_lib.SchemaHelper:
     raise NotImplementedError(type(self))
 
   def get_data_slice(
@@ -46,6 +46,15 @@ def get_data_slice(
   def get_schema(self) -> kd.types.SchemaItem:
     return self._get_schema_helper().get_schema()
 
+  def get_schema_at(
+      self, path: data_slice_path_lib.DataSlicePath
+  ) -> kd.types.SchemaItem:
+    """Returns the schema of the DataSlice at the given path."""
+    schema_helper = self._get_schema_helper()
+    return schema_helper.get_subschema_at(
+        schema_helper.get_schema_node_name_for_data_slice_path(path)
+    )
+
   def exists(self, path: data_slice_path_lib.DataSlicePath) -> bool:
     return self._get_schema_helper().exists(path)