corvic-engine 0.3.0rc61__cp38-abi3-win_amd64.whl → 0.3.0rc63__cp38-abi3-win_amd64.whl

corvic/embed/node2vec.py

@@ -39,6 +39,8 @@ class KeyedVectors:
           index_to_key: mapping of index to key struct
           key_field_order: order of key struct fields used for index operations
         """
+        if dim <= 0:
+            raise InvalidArgumentError("number of dimensions must be positive")
         self.dim = dim
         self._index_to_key = index_to_key
         self._key_field_order = key_field_order

corvic/op_graph/_transformations.py

@@ -73,7 +73,7 @@ def _replace_concat_op_source(
     return op.concat(new_tables, concat_op.how)
 
 
-def replace_op_source(  # noqa: C901
+def replace_op_source(  # noqa: C901, PLR0915
     root_op: op.Op, source_to_replace: op.Op, new_source: op.Op
 ) -> Ok[op.Op] | InvalidArgumentError:
     for source in root_op.sources():
@@ -98,6 +98,8 @@ def replace_op_source(  # noqa: C901
             return new_source.select_columns(root_op.columns)
         case op.LimitRows():
             return new_source.limit_rows(root_op.num_rows)
+        case op.OffsetRows():
+            return new_source.offset_rows(root_op.num_rows)
         case op.OrderBy():
             return new_source.order_by(root_op.columns, desc=root_op.desc)
         case op.FilterRows():

corvic/op_graph/ops.py

@@ -71,6 +71,7 @@ ProtoOp = (
     | table_pb2.JoinOp
     | table_pb2.SelectColumnsOp
     | table_pb2.LimitRowsOp
+    | table_pb2.OffsetRowsOp
     | table_pb2.OrderByOp
     | table_pb2.FilterRowsOp
     | table_pb2.DistinctRowsOp
@@ -162,6 +163,15 @@ def from_proto(
 ) -> LimitRows: ...
 
 
+@overload
+def from_proto(
+    proto: table_pb2.OffsetRowsOp,
+    parent_ops: list[Op] | None = None,
+    *,
+    skip_validate: bool = False,
+) -> OffsetRows: ...
+
+
 @overload
 def from_proto(
     proto: table_pb2.OrderByOp,
@@ -490,6 +500,8 @@ def from_proto(  # noqa: C901, PLR0915
             proto = table_pb2.TableComputeOp(select_columns=proto)
         case table_pb2.LimitRowsOp():
             proto = table_pb2.TableComputeOp(limit_rows=proto)
+        case table_pb2.OffsetRowsOp():
+            proto = table_pb2.TableComputeOp(offset_rows=proto)
         case table_pb2.OrderByOp():
             proto = table_pb2.TableComputeOp(order_by=proto)
         case table_pb2.FilterRowsOp():
@@ -992,6 +1004,12 @@ class _Base(OneofProtoWrapper[table_pb2.TableComputeOp], ABC):
             proto = table_pb2.LimitRowsOp(source=self._proto, num_rows=num_rows)
         return Ok(from_proto(proto, skip_validate=True))
 
+    def offset_rows(self, num_rows: int) -> InvalidArgumentError | Ok[OffsetRows]:
+        if num_rows <= 0:
+            return InvalidArgumentError("num_rows must be positive")
+        proto = table_pb2.OffsetRowsOp(source=self._proto, num_rows=num_rows)
+        return Ok(from_proto(proto, skip_validate=True))
+
     def order_by(
         self, columns: Sequence[str], *, desc: bool
     ) -> InvalidArgumentError | Ok[OrderBy]:
@@ -1903,6 +1921,23 @@ class LimitRows(_Base):
         return [self.source]
 
 
+class OffsetRows(_Base):
+    """Limit the number of rows in a table."""
+
+    @property
+    def num_rows(self) -> int:
+        return self._proto.offset_rows.num_rows
+
+    @property
+    def source(self) -> Op:
+        if self._parents:
+            return self._parents[0]
+        return from_proto(self._proto.offset_rows.source, skip_validate=True)
+
+    def sources(self):
+        return [self.source]
+
+
 class OrderBy(_Base):
     """Order the rows in a table."""
 
@@ -2756,6 +2791,7 @@ Op = (
     | Join
     | SelectColumns
     | LimitRows
+    | OffsetRows
     | OrderBy
     | FilterRows
     | DistinctRows
@@ -2798,6 +2834,7 @@ _COMPUTE_OP_FIELD_NAME_TO_OP: Final = {
     "join": Join,
     "select_columns": SelectColumns,
     "limit_rows": LimitRows,
+    "offset_rows": OffsetRows,
     "order_by": OrderBy,
     "filter_rows": FilterRows,
     "distinct_rows": DistinctRows,
@@ -3323,6 +3360,7 @@ class Schema(Sequence[Field]):
 
             case (
                 LimitRows()
+                | OffsetRows()
                 | OrderBy()
                 | FilterRows()
                 | DistinctRows()

corvic/sql/parse_ops.py

@@ -24,6 +24,7 @@ _SqlComputableOp = (
     | op_graph.op.Join
     | op_graph.op.SelectColumns
     | op_graph.op.LimitRows
+    | op_graph.op.OffsetRows
     | op_graph.op.OrderBy
     | op_graph.op.FilterRows
     | op_graph.op.DistinctRows
@@ -74,6 +75,7 @@ def can_be_sql_computed(
             | op_graph.op.Join()
             | op_graph.op.SelectColumns()
             | op_graph.op.LimitRows()
+            | op_graph.op.OffsetRows()
             | op_graph.op.OrderBy()
             | op_graph.op.FilterRows()
             | op_graph.op.DistinctRows()
@@ -362,6 +364,12 @@ class _OpLogParser:
     ) -> Ok[sqlglot.exp.Query] | InvalidArgumentError | NoRowsError:
         return self.parse(op.source).map(lambda query: query.limit(op.num_rows))
 
+    def _offset_rows_to_sql(
+        self,
+        op: op_graph.op.OffsetRows,
+    ) -> Ok[sqlglot.exp.Query] | InvalidArgumentError | NoRowsError:
+        return self.parse(op.source).map(lambda query: query.offset(op.num_rows))
+
     def _order_by_to_sql(
         self,
         op: op_graph.op.OrderBy,
@@ -715,6 +723,8 @@ class _OpLogParser:
                 return self._select_columns_to_sql(op)
             case op_graph.op.LimitRows():
                 return self._limit_rows_to_sql(op)
+            case op_graph.op.OffsetRows():
+                return self._offset_rows_to_sql(op)
             case op_graph.op.OrderBy():
                 return self._order_by_to_sql(op)
             case op_graph.op.FilterRows():

corvic/system/_embedder.py

@@ -1,6 +1,6 @@
 import dataclasses
 from collections.abc import Sequence
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 import polars as pl
@@ -9,6 +9,12 @@ from typing_extensions import Protocol
 from corvic import orm
 from corvic.result import InternalError, InvalidArgumentError, Ok
 
+if TYPE_CHECKING:
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
+
 
 @dataclasses.dataclass
 class EmbedTextContext:
@@ -64,6 +70,12 @@ class ImageEmbedder(Protocol):
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError: ...
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class ClipText(TextEmbedder):
     """Clip Text embedder.
 
@@ -76,28 +88,39 @@ class ClipText(TextEmbedder):
     overcoming several major challenges in computer vision.
     """
 
-    def embed(
-        self, context: EmbedTextContext
-    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
-        import torch
+    def _load_models(self):
         from transformers import (
             CLIPModel,
             CLIPProcessor,
         )
 
         model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
         )
         processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
         )
-        model.eval()
+        return ClipModels(model=model, processor=processor)
+
+    def embed(
+        self, context: EmbedTextContext
+    ) -> Ok[EmbedTextResult] | InvalidArgumentError | InternalError:
         match context.expected_coordinate_bitwidth:
             case 64:
                 coord_dtype = pl.Float64()
             case 32:
                 coord_dtype = pl.Float32()
 
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
+        model.eval()
+
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.Tensor] = processor(  # pyright: ignore[reportAssignmentType]
                 text=context.inputs,

corvic/system/_image_embedder.py

@@ -1,3 +1,4 @@
+import dataclasses
 from io import BytesIO
 from typing import TYPE_CHECKING, Any
 
@@ -13,6 +14,10 @@ from corvic.system._embedder import (
 
 if TYPE_CHECKING:
     from PIL import Image
+    from transformers import (
+        CLIPModel,
+        CLIPProcessor,
+    )
 
 
 class RandomImageEmbedder(ImageEmbedder):
@@ -58,6 +63,12 @@ def image_from_bytes(
         return InvalidArgumentError("invalid image format")
 
 
+@dataclasses.dataclass
+class ClipModels:
+    model: "CLIPModel"
+    processor: "CLIPProcessor"
+
+
 class Clip(ImageEmbedder):
     """Clip image embedder.
 
@@ -70,6 +81,23 @@ class Clip(ImageEmbedder):
     overcoming several major challenges in computer vision.
     """
 
+    def _load_models(self):
+        from transformers import (
+            CLIPModel,
+            CLIPProcessor,
+        )
+
+        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+        )
+        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
+            pretrained_model_name_or_path="openai/clip-vit-base-patch32",
+            revision="5812e510083bb2d23fa43778a39ac065d205ed4d",
+            use_fast=False,
+        )
+        return ClipModels(model=model, processor=processor)
+
     def embed(
         self, context: EmbedImageContext
     ) -> Ok[EmbedImageResult] | InvalidArgumentError | InternalError:
@@ -99,20 +127,13 @@ class Clip(ImageEmbedder):
                 )
             )
 
-        import torch
-        from transformers import (
-            CLIPModel,
-            CLIPProcessor,
-        )
-
-        model: CLIPModel = CLIPModel.from_pretrained(  # pyright: ignore[reportUnknownMemberType]
-            "openai/clip-vit-base-patch32"
-        )
-        processor: CLIPProcessor = CLIPProcessor.from_pretrained(  # pyright: ignore[reportUnknownMemberType, reportAssignmentType]
-            "openai/clip-vit-base-patch32"
-        )
+        models = self._load_models()
+        model = models.model
+        processor = models.processor
         model.eval()
 
+        import torch
+
         with torch.no_grad():
             inputs: dict[str, torch.FloatTensor] = processor(  # pyright: ignore[reportAssignmentType]
                 images=images, return_tensors="pt"

corvic/system/_planner.py

@@ -166,6 +166,9 @@ class OpGraphPlanner:
             case op_graph.op.LimitRows() | op_graph.op.SampleRows():
                 source_rows = cls.count_rows_upperbound(op.source)
                 num_rows = min(op.num_rows, source_rows)
+            case op_graph.op.OffsetRows():
+                source_rows = cls.count_rows_upperbound(op.source)
+                num_rows = max(source_rows - op.num_rows, 0)
             case op_graph.op.Empty():
                 num_rows = 0
             case op_graph.op.AggregateColumns():