pyspiral 0.2.5__cp310-abi3-macosx_11_0_arm64.whl → 0.3.1__cp310-abi3-macosx_11_0_arm64.whl

spiral/substrait_.py

@@ -251,7 +251,7 @@ class SubstraitConverter:
 
                 if field_name in self.key_names:
                     # This is a key column, so we need to select it from the scope.
-                    return se.var(field_name)
+                    return se.key(field_name)
 
                 scope = se.getitem(scope, field_name)
                 scope_type = scope_type.field(ref.field).type

spiral/table.py

@@ -3,11 +3,13 @@ from typing import TYPE_CHECKING, Literal
 
 import pyarrow as pa
 
-from spiral import expressions as se
-from spiral.config import FILE_FORMAT, Config
 from spiral.core.core import Table as CoreTable
-from spiral.core.core import flush_wal, write
+from spiral.core.core import TableMaintenance, TableTransaction
+from spiral.core.spec import Schema
 from spiral.expressions.base import Expr, ExprLike
+from spiral.maintenance import Maintenance
+from spiral.settings import FILE_FORMAT
+from spiral.txn import Transaction
 
 if TYPE_CHECKING:
     import duckdb
@@ -53,7 +55,7 @@ class Table(Expr):
         from spiral import expressions as se
 
         if item in self._key_columns:
-            return se.var(name=item)
+            return se.key(name=item)
 
         return super().__getitem__(item)
 
@@ -72,7 +74,7 @@ class Table(Expr):
 
         from spiral import expressions as se
 
-        return se.merge(se.pack({key: se.var(key) for key in key_paths}), super().select(*other_paths, exclude=exclude))
+        return se.merge(se.pack({key: se.key(key) for key in key_paths}), super().select(*other_paths, exclude=exclude))
 
     @property
     def key_schema(self) -> pa.Schema:
@@ -80,7 +82,7 @@ class Table(Expr):
         return self._key_schema
 
     @property
-    def schema(self) -> pa.Schema:
+    def schema(self) -> Schema:
         """Returns the FULL schema of the table.
 
         NOTE: This can be expensive for large tables.
@@ -111,8 +113,6 @@ class Table(Expr):
         where: ExprLike | None = None,
         asof: datetime | int | str = None,
         exclude_keys: bool = False,
-        # TODO(marko): Support config.
-        # config: Config | None = None,
     ) -> "Scan":
         """Reads the table. If projections are not provided, the entire table is read.
 
@@ -128,7 +128,6 @@ class Table(Expr):
             where=where,
             asof=asof,
             exclude_keys=exclude_keys,
-            # config=config,
         )
 
     # NOTE: "vortex" is valid format. We don't want that visible in the API docs.
@@ -137,21 +136,36 @@ class Table(Expr):
         expr: ExprLike,
         *,
         format: Literal["parquet"] | None = None,
-        # TODO(joe): support group_by, and config
-        config: Config | None = None,
+        partition_size_bytes: int | None = None,
     ) -> None:
         """Write an item to the table inside a single transaction.
 
         :param expr: The expression to write. Must evaluate to a struct array.
         :param format: the format to write the data in. Defaults to "parquet".
-        :param config: The configuration to use for this write.
+        :param partition_size_bytes: The maximum partition size in bytes.
         """
-        write(
-            self._table,
-            se.lift(expr).__expr__,
-            format=format or FILE_FORMAT,
-            partition_size=config.partition_file_min_size if config else None,
-        )
-        # Flush the WAL if configured.
-        if config is not None and config.flush_wal_on_write:
-            flush_wal(self._table, manifest_format=format or FILE_FORMAT)
+        format = format or FILE_FORMAT
+
+        with self.txn(format=format) as txn:
+            txn.write(
+                expr,
+                partition_size_bytes=partition_size_bytes,
+            )
+
+    # NOTE: "vortex" is valid format. We don't want that visible in the API docs.
+    def txn(self, format: Literal["parquet"] | None = None) -> Transaction:
+        """Begins a new transaction. Transaction must be committed for writes to become visible.
+
+        IMPORTANT: While transaction can be used to atomically write data to the table,
+             it is important that the primary key columns are unique within the transaction.
+
+        :param format: The format to use for the transaction. Defaults to "parquet".
+        """
+        return Transaction(TableTransaction(self._table.metastore, format or FILE_FORMAT))
+
+    def maintenance(self, format: Literal["parquet"] | None = None) -> Maintenance:
+        """Maintenance tasks for the table.
+
+        :param format: The format to use. Defaults to "parquet".
+        """
+        return Maintenance(TableMaintenance(self._table.metastore, format or FILE_FORMAT))

spiral/txn.py

@@ -0,0 +1,48 @@
+from spiral.core.core import TableTransaction
+from spiral.expressions import ExprLike
+
+
+class Transaction:
+    """Spiral table transaction.
+
+    IMPORTANT: While transaction can be used to atomically write data to the table,
+            it is important that the primary key columns are unique within the transaction.
+    """
+
+    def __init__(self, transaction: TableTransaction):
+        self._transaction = transaction
+
+    @property
+    def status(self) -> str:
+        """The status of the transaction."""
+        return self._transaction.status
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            self._transaction.commit()
+        else:
+            self._transaction.abort()
+
+    def write(self, expr: ExprLike, *, partition_size_bytes: int | None = None):
+        """Write an item to the table inside a single transaction.
+
+        :param expr: The expression to write. Must evaluate to a struct array.
+        :param partition_size_bytes: The maximum partition size in bytes.
+            If not provided, the default partition size is used.
+        """
+        from spiral import expressions as se
+
+        expr = se.lift(expr)
+
+        self._transaction.write(expr.__expr__, partition_size_bytes=partition_size_bytes)
+
+    def commit(self):
+        """Commit the transaction."""
+        self._transaction.commit()
+
+    def abort(self):
+        """Abort the transaction."""
+        self._transaction.abort()

spiral/config.py

@@ -1,26 +0,0 @@
-import os
-
-from pydantic_settings import BaseSettings, SettingsConfigDict
-
-FILE_FORMAT = os.environ.get("SPIRAL_FILE_FORMAT", "parquet")
-
-
-class Config(BaseSettings):
-    model_config = SettingsConfigDict(
-        env_nested_delimiter="__",
-        env_prefix="SPIRAL_CORE__",
-        frozen=True,
-    )
-
-    partition_file_min_size: int = 256 * 1024 * 1024  # 256MB
-    flush_wal_on_write: bool = False
-
-    # TODO(marko): Support config. Unused after migration to Rust.
-    # #: Defaults to ThreadPoolExecutor's default (based on os.cpu_count().
-    # scan_num_threads: int | None = 61  # 61 is used by Golang and Tokio, for some reason...
-    #
-    # #: The duration of WAL that is preserved to allow for txn conflict resolution.
-    # transaction_window: int = 0 if DEV else timedelta(days=1).total_seconds()
-    #
-    # #: Truncation length of string statistics.
-    # string_truncation_length: int = 1024