Batch events

Author: blythed

plugins/snowflake/superduper_snowflake/data_backend.py

@@ -12,6 +12,7 @@
 from snowflake.snowpark.types import BooleanType, StringType, VariantType
 from superduper import CFG, logging
 from superduper.backends.base.data_backend import BaseDataBackend
+from superduper.base.event import Create, CreateTable
 from superduper.base.query import Query
 from superduper.base.schema import Schema
 from watchdog.events import FileSystemEventHandler
@@ -115,6 +116,13 @@ def db(self, value):
         """
         self._db = value
 
+    def create_tables_and_schemas(self, events: t.List[CreateTable]):
+        """Create tables and schemas in the data-backend.
+
+        :param events: The events to create.
+        """
+        raise NotImplementedError
+
     def create_table_and_schema(self, identifier: str, schema: Schema, primary_id: str):
         """Create a schema in the data-backend.
 

superduper/backends/base/backends.py

@@ -45,12 +45,13 @@ def get_tool(self, uuid: str):
         tool_id = self.uuid_tool_mapping[uuid]
         return self.tools[tool_id]
 
-    def put_component(self, component: 'Component', **kwargs):
+    def put_component(self, component: 'Component', uuid: str):
         """Put a component to the backend.
 
         :param component: Component to put.
         :param kwargs: kwargs dictionary.
         """
+        component = self.db.load(component=component, uuid=uuid)
         logging.info(
             f'Putting component: {component.huuid} on to {self.__class__.__name__}'
         )
@@ -68,7 +69,7 @@ def put_component(self, component: 'Component', **kwargs):
             return
         self.tool_uuid_mapping[tool.identifier].add(component.uuid)
         self.tools[tool.identifier] = tool
-        tool.initialize(**kwargs)
+        tool.initialize()
 
     def drop_component(self, component: str, identifier: str):
         """Drop the component from backend.

superduper/backends/base/compute.py

@@ -65,7 +65,7 @@ def initialize(self):
                     self.put_component(m)
 
     @abstractmethod
-    def put_component(self, component: 'Component'):
+    def put_component(self, component: 'Component', uuid: str):
         """Create handler on component declare.
 
         :param component: Component to put.

superduper/backends/base/data_backend.py

@@ -12,6 +12,7 @@
 
 if t.TYPE_CHECKING:
     from superduper.base.schema import Schema
+    from superduper.base.event import CreateTable
 
 
 class BaseDataBackend(ABC):
@@ -85,6 +86,19 @@ def db(self, value):
         """
         self._db = value
 
+    def create_tables_and_schemas(self, events: t.List['CreateTable']):
+        """Create a schema in the data-backend.
+
+        :param events: List of `CreateTable` events.
+        """
+        from superduper.base.schema import Schema
+        for event in events:
+            self.create_table_and_schema(
+                event.identifier,
+                schema=Schema.build(**event.fields),
+                primary_id=event.primary_id,
+            )
+
     @abstractmethod
     def create_table_and_schema(
         self, identifier: str, schema: 'Schema', primary_id: str

superduper/backends/base/scheduler.py

@@ -9,6 +9,7 @@
 from superduper import CFG, logging
 from superduper.backends.base.backends import BaseBackend
 from superduper.base.base import Base
+from superduper.base.event import Create, CreateTable, PutComponent, Update
 
 DependencyType = t.Union[t.Dict[str, str], t.Sequence[t.Dict[str, str]]]
 
@@ -140,6 +141,30 @@ def _consume_event_type(event_type, ids, table, db: 'Datalayer'):
     db.cluster.compute.release_futures(context)
 
 
+def cluster_events(events: t.List[Base]) -> t.Tuple[t.List['CreateTable'], t.List['Create'], t.List['Job']]:
+    """
+    Cluster events into table, create and job events.
+
+    :param events: List of events to be clustered.
+    :return: Tuple of table events, create events and job events.
+    """
+    from superduper.base.metadata import Job
+    table_events = []
+    create_events = []
+    job_events = []
+    put_events = []
+    for event in events:
+        if isinstance(event, CreateTable):
+            table_events.append(event)
+        elif isinstance(event, (Update, Create)):
+            create_events.append(event)
+        elif isinstance(event, Job):
+            job_events.append(event)
+        elif isinstance(event, PutComponent):
+            put_events.append(event)
+    return table_events, create_events, put_events, job_events
+
+
 def consume_events(events: t.List[Base], table: str, db: 'Datalayer'):
     """
     Consume events from table queue.
@@ -152,7 +177,32 @@ def consume_events(events: t.List[Base], table: str, db: 'Datalayer'):
         logging.info(f'Consuming {len(events)} events on {table}.')
         consume_streaming_events(events=events, table=table, db=db)
     else:
-        logging.info(f'Consuming {len(events)} _apply events')
-        for event in events:
-            event.execute(db)
+        table_events, create_events, put_events, job_events = cluster_events(events)
+
+        if table_events:
+            logging.info(f'Consuming {len(events)} `CreateTable` events')
+            CreateTable.batch_execute(
+                events=table_events,
+                db=db,
+            )
+
+        if create_events:
+            logging.info(f'Consuming {len(events)} `Create` events')
+            Create.batch_execute(
+                events=create_events,
+                db=db,
+            )
+
+        if put_events:
+            logging.info(f'Consuming {len(events)} `PutComponent` events')
+            PutComponent.batch_execute(
+                events=put_events,
+                db=db,
+            )
+
+        if job_events:
+            logging.info(f'Consuming {len(events)} jobs (`Job`)')
+            for job in job_events:
+                job.execute(db)
+
         return

superduper/backends/local/cdc.py

@@ -32,9 +32,8 @@ def list_uuids(self):
         """List UUIDs of components."""
         return list(self._trigger_uuid_mapping.values())
 
-    def put_component(self, component):
-        assert isinstance(component, CDC)
-        self.triggers.add((component.component, component.identifier))
+    def put_component(self, component: str, uuid: str):
+        self.triggers.add((component, uuid))
 
     def drop_component(self, component, identifier):
         c = self.db.load(component=component, identifier=identifier)

superduper/backends/local/compute.py

@@ -77,7 +77,7 @@ def list_uuids(self):
     def drop_component(self, component: str, identifier: str):
         """Drop a component from the compute."""
 
-    def put_component(self, component: 'Component'):
+    def put_component(self, component: 'Component', uuid: str):
         """Create a handler on compute."""
 
     def initialize(self):

superduper/base/apply.py

@@ -8,7 +8,7 @@
 from superduper import CFG, Component, logging
 from superduper.base import exceptions
 from superduper.base.document import Document
-from superduper.base.event import Create, Signal, Update
+from superduper.base.event import Create, PutComponent, Signal, Update
 from superduper.components.component import running_status
 from superduper.misc.tree import dict_to_tree
 
@@ -62,9 +62,10 @@ def apply(
         force = db.cfg.force_apply
 
     # This holds a record of the changes
-    diff: t.Dict = {}
+    diff = {}
+
     # context allows us to track the origin of the component creation
-    create_events, job_events = _apply(
+    table_events, create_events, put_events, job_events = _apply(
         db=db,
         object=object,
         context=object.uuid,
@@ -87,6 +88,14 @@ def apply(
 
     logging.info('Found these changes and/ or additions that need to be made:')
 
+    logging.info('-' * 100)
+    logging.info('TABLE EVENTS:')
+    logging.info('-' * 100)
+    steps = {t.huuid: str(i) for i, t in enumerate(table_events.values())}
+
+    for i, t in enumerate(table_events.values()):
+        logging.info(f'[{i}]: {t.huuid}')
+
     logging.info('-' * 100)
     logging.info('METADATA EVENTS:')
     logging.info('-' * 100)
@@ -104,6 +113,14 @@ def apply(
         else:
             logging.info(f'[{i}]: {c.huuid}: {c.__class__.__name__}')
 
+    logging.info('-' * 100)
+    logging.info('PUT EVENTS:')
+    logging.info('-' * 100)
+    steps = {p.huuid: str(i) for i, p in enumerate(put_events.values())}
+
+    for i, t in enumerate(put_events.values()):
+        logging.info(f'[{i}]: {t.huuid}')
+
     logging.info('-' * 100)
     logging.info('JOBS EVENTS:')
     logging.info('-' * 100)
@@ -124,7 +141,9 @@ def apply(
     logging.info('-' * 100)
 
     events = [
+        *list(table_events.values()),
         *list(create_events.values()),
+        *list(put_events.values()),
         *list(job_events.values()),
         Signal(context=object.uuid, msg='done'),
     ]
@@ -192,12 +211,14 @@ def _apply(
     object.db = db
 
     create_events = {}
+    table_events = {}
+    put_events = {}
     children = []
 
     def wrapper(child):
         nonlocal create_events
 
-        c, j = _apply(
+        t, c, p, j = _apply(
             db=db,
             object=child,
             context=context,
@@ -209,6 +230,8 @@ def wrapper(child):
 
         job_events.update(j)
         create_events.update(c)
+        table_events.update(t)
+        put_events.update(p)
         children.append((child.component, child.identifier, child.uuid))
         return f'&:component:{child.huuid}'
 
@@ -252,7 +275,7 @@ def wrapper(child):
     serialized = db._save_artifact(serialized.encode())
 
     if apply_status == 'same':
-        return create_events, job_events
+        return table_events, create_events, put_events, job_events
 
     elif apply_status == 'new':
 
@@ -264,11 +287,23 @@ def wrapper(child):
             children=children,
         )
 
+        table_events.update(object.create_table_events())
+
         these_job_events = object.create_jobs(
             event_type='apply',
             jobs=list(job_events.values()),
             context=context,
         )
+
+        for service in object.services:
+            put_events[f'{object.huuid}/{service}'] = PutComponent(
+                component=object.component,
+                identifier=object.identifier,
+                uuid=object.uuid,
+                context=context,
+                service=service,
+            )
+
     elif apply_status == 'breaking':
 
         metadata_event = Create(
@@ -278,12 +313,23 @@ def wrapper(child):
             parent=parent,
             children=children,
         )
+        
+        table_events.update(object.create_table_events())
 
         these_job_events = object.create_jobs(
             event_type='apply',
             jobs=list(job_events.values()),
             context=context,
         )
+
+        for service in object.services:
+            put_events[f'{object.huuid}/{service}'] = PutComponent(
+                component=object.component,
+                identifier=object.identifier,
+                uuid=object.uuid,
+                context=context,
+                service=service,
+            )
     else:
         assert apply_status == 'update'
 
@@ -318,4 +364,4 @@ def wrapper(child):
 
     create_events[metadata_event.huuid] = metadata_event
     job_events.update({jj.huuid: jj for jj in these_job_events})
-    return create_events, job_events
+    return table_events, create_events, put_events, job_events