pixeltable 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

pixeltable/metadata/converters/convert_29.py

@@ -80,7 +80,7 @@ def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], A
                         rolled_kwargs = kwargs.pop(param['name'])
 
         if rolled_args is not None:
-            assert rolled_args['_classname'] in {'InlineArray', 'InlineList'}
+            assert rolled_args['_classname'] in ('InlineArray', 'InlineList')
             new_args.extend(rolled_args['components'])
         if rolled_kwargs is not None:
             assert rolled_kwargs['_classname'] == 'InlineDict'

pixeltable/metadata/converters/convert_30.py

@@ -0,0 +1,50 @@
+import copy
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import (
+    convert_table_record,
+    convert_table_schema_version_record,
+    convert_table_version_record,
+)
+from pixeltable.metadata.schema import Table, TableSchemaVersion, TableVersion
+
+
+@register_converter(version=30)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_record(engine, table_record_updater=__update_table_record)
+    convert_table_version_record(engine, table_version_record_updater=__update_table_version_record)
+    convert_table_schema_version_record(
+        engine, table_schema_version_record_updater=__update_table_schema_version_record
+    )
+
+
+def __update_table_record(record: Table) -> None:
+    """
+    Update TableMd with table_id
+    """
+    assert isinstance(record.md, dict)
+    md = copy.copy(record.md)
+    md['tbl_id'] = str(record.id)
+    record.md = md
+
+
+def __update_table_version_record(record: TableVersion) -> None:
+    """
+    Update TableVersion with table_id.
+    """
+    assert isinstance(record.md, dict)
+    md = copy.copy(record.md)
+    md['tbl_id'] = str(record.tbl_id)
+    record.md = md
+
+
+def __update_table_schema_version_record(record: TableSchemaVersion) -> None:
+    """
+    Update TableSchemaVersion with table_id.
+    """
+    assert isinstance(record.md, dict)
+    md = copy.copy(record.md)
+    md['tbl_id'] = str(record.tbl_id)
+    record.md = md

pixeltable/metadata/converters/util.py

@@ -5,7 +5,7 @@ from uuid import UUID
 
 import sqlalchemy as sql
 
-from pixeltable.metadata.schema import Function, Table, TableSchemaVersion
+from pixeltable.metadata.schema import Function, Table, TableSchemaVersion, TableVersion
 
 __logger = logging.getLogger('pixeltable')
 
@@ -143,3 +143,28 @@ def __update_schema_column(table_schema_version_md: dict, schema_column_updater:
     assert isinstance(cols, dict)
     for schema_col in cols.values():
         schema_column_updater(schema_col)
+
+
+def convert_table_record(engine: sql.engine.Engine, table_record_updater: Optional[Callable[[Table], None]]) -> None:
+    with sql.orm.Session(engine, future=True) as session:
+        for record in session.query(Table).all():
+            table_record_updater(record)
+        session.commit()
+
+
+def convert_table_version_record(
+    engine: sql.engine.Engine, table_version_record_updater: Optional[Callable[[TableVersion], None]]
+) -> None:
+    with sql.orm.Session(engine, future=True) as session:
+        for record in session.query(TableVersion).all():
+            table_version_record_updater(record)
+        session.commit()
+
+
+def convert_table_schema_version_record(
+    engine: sql.engine.Engine, table_schema_version_record_updater: Optional[Callable[[TableSchemaVersion], None]]
+) -> None:
+    with sql.orm.Session(engine, future=True) as session:
+        for record in session.query(TableSchemaVersion).all():
+            table_schema_version_record_updater(record)
+        session.commit()

pixeltable/metadata/notes.py

@@ -2,6 +2,7 @@
 # rather than as a comment, so that the existence of a description can be enforced by
 # the unit tests when new versions are added.
 VERSION_NOTES = {
+    31: 'Add table ids to metadata structs',
     30: 'Store default values and constant arguments as literals',
     29: 'Add user and additional_md fields to metadata structs',
     28: 'Enable view creation from DataFrame with select clause',

pixeltable/metadata/schema.py

@@ -153,6 +153,7 @@ class ViewMd:
 
 @dataclasses.dataclass
 class TableMd:
+    tbl_id: str  # uuid.UUID
     name: str
 
     user: Optional[str]
@@ -199,6 +200,7 @@ class Table(Base):
 
 @dataclasses.dataclass
 class TableVersionMd:
+    tbl_id: str  # uuid.UUID
     created_at: float  # time.time()
     version: int
     schema_version: int
@@ -234,6 +236,7 @@ class TableSchemaVersionMd:
     Records all versioned table metadata.
     """
 
+    tbl_id: str  # uuid.UUID
     schema_version: int
     preceding_schema_version: Optional[int]
     columns: dict[int, SchemaColumn]  # col_id -> SchemaColumn

pixeltable/store.py

@@ -99,9 +99,9 @@ class StoreBase:
 
         # v_min/v_max indices: speeds up base table scans needed to propagate a base table insert or delete
         idx_name = f'vmin_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_min_col, postgresql_using=Env.get().dbms.version_index_type))
         idx_name = f'vmax_idx_{self.tbl_version.id.hex}'
-        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using='brin'))
+        idxs.append(sql.Index(idx_name, self.v_max_col, postgresql_using=Env.get().dbms.version_index_type))
 
         self.sa_tbl = sql.Table(self._storage_name(), self.sa_md, *all_cols, *idxs)
 

pixeltable/type_system.py

@@ -512,7 +512,7 @@ class StringType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.STRING, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -602,7 +602,7 @@ class TimestampType(ColumnType):
     def __init__(self, nullable: bool = False):
         super().__init__(self.Type.TIMESTAMP, nullable=nullable)
 
-    def has_supertype(self):
+    def has_supertype(self) -> bool:
         return not self.nullable
 
     @classmethod
@@ -768,7 +768,7 @@ class JsonType(ColumnType):
         a_type = a.get('type')
         b_type = b.get('type')
 
-        if a_type in {'string', 'integer', 'number', 'boolean', 'object', 'array'} and a_type == b_type:
+        if a_type in ('string', 'integer', 'number', 'boolean', 'object', 'array') and a_type == b_type:
             # a and b both have the same type designation, but are not identical. This can happen if
             # (for example) they have validators or other attributes that differ. In this case, we
             # generalize to {'type': t}, where t is their shared type, with no other qualifications.
@@ -1170,6 +1170,20 @@ class DocumentType(ColumnType):
         XML = 3
         TXT = 4
 
+        @classmethod
+        def from_extension(cls, ext: str) -> Optional['DocumentType.DocumentFormat']:
+            if ext in ('.htm', '.html'):
+                return cls.HTML
+            if ext == '.md':
+                return cls.MD
+            if ext == '.pdf':
+                return cls.PDF
+            if ext == '.xml':
+                return cls.XML
+            if ext == '.txt':
+                return cls.TXT
+            return None
+
     def __init__(self, nullable: bool = False, doc_formats: Optional[str] = None):
         super().__init__(self.Type.DOCUMENT, nullable=nullable)
         self.doc_formats = doc_formats
@@ -1203,9 +1217,7 @@ class DocumentType(ColumnType):
         assert isinstance(val, str)
         from pixeltable.utils.documents import get_document_handle
 
-        dh = get_document_handle(val)
-        if dh is None:
-            raise excs.Error(f'Not a recognized document format: {val}')
+        _ = get_document_handle(val)
 
 
 T = typing.TypeVar('T')
@@ -1240,7 +1252,7 @@ class _PxtType:
     `ColumnType`.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         raise TypeError(f'Type `{type(self)}` cannot be instantiated.')
 
     @classmethod

pixeltable/utils/arrow.py

@@ -11,14 +11,19 @@ PA_TO_PXT_TYPES: dict[pa.DataType, ts.ColumnType] = {
     pa.large_string(): ts.StringType(nullable=True),
     pa.timestamp('us', tz=datetime.timezone.utc): ts.TimestampType(nullable=True),
     pa.bool_(): ts.BoolType(nullable=True),
-    pa.uint8(): ts.IntType(nullable=True),
     pa.int8(): ts.IntType(nullable=True),
-    pa.uint32(): ts.IntType(nullable=True),
-    pa.uint64(): ts.IntType(nullable=True),
+    pa.int16(): ts.IntType(nullable=True),
     pa.int32(): ts.IntType(nullable=True),
     pa.int64(): ts.IntType(nullable=True),
+    pa.uint8(): ts.IntType(nullable=True),
+    pa.uint16(): ts.IntType(nullable=True),
+    pa.uint32(): ts.IntType(nullable=True),
+    pa.uint64(): ts.IntType(nullable=True),
     pa.float32(): ts.FloatType(nullable=True),
     pa.float64(): ts.FloatType(nullable=True),
+    pa.date32(): ts.StringType(nullable=True),  # date32 is not supported in pixeltable, use string
+    pa.date64(): ts.StringType(nullable=True),  # date64 is not supported in pixeltable, use string
+    pa.binary(): None,  # cannot import binary (inline image)
 }
 
 PXT_TO_PA_TYPES: dict[type[ts.ColumnType], pa.DataType] = {
@@ -43,7 +48,7 @@ def to_pixeltable_type(arrow_type: pa.DataType, nullable: bool) -> Optional[ts.C
         return ts.TimestampType(nullable=nullable)
     elif arrow_type in PA_TO_PXT_TYPES:
         pt = PA_TO_PXT_TYPES[arrow_type]
-        return pt.copy(nullable=nullable)
+        return pt.copy(nullable=nullable) if pt is not None else None
     elif isinstance(arrow_type, pa.FixedShapeTensorType):
         dtype = to_pixeltable_type(arrow_type.value_type, nullable)
         if dtype is None:
@@ -111,6 +116,28 @@ def iter_tuples(batch: Union[pa.Table, pa.RecordBatch]) -> Iterator[dict[str, An
         yield {col_name: values[i] for col_name, values in pydict.items()}
 
 
+def _ar_val_to_pxt_val(val: Any, pxt_type: ts.ColumnType) -> Any:
+    """Convert a value to insertable format"""
+    if val is None:
+        return None
+    if pxt_type.is_float_type():
+        return float(val)
+    elif pxt_type.is_int_type():
+        return int(val)
+    elif pxt_type.is_bool_type():
+        return bool(val)
+    elif pxt_type.is_string_type():
+        return str(val)
+    elif pxt_type.is_timestamp_type():
+        if isinstance(val, str):
+            return datetime.datetime.fromisoformat(val)
+        if isinstance(val, datetime.datetime):
+            return val
+    elif pxt_type.is_array_type():
+        return pxt_type.create_literal(val)
+    raise ValueError(f'Unsupported type {pxt_type} for value {val}')
+
+
 def iter_tuples2(
     batch: Union[pa.Table, pa.RecordBatch], col_mapping: Optional[dict[str, str]], schema: dict[str, ts.ColumnType]
 ) -> Iterator[dict[str, Any]]:
@@ -124,8 +151,6 @@ def iter_tuples2(
     for i in range(batch_size):
         # Convert a row to insertable format
         yield {
-            (pxt_name := col_name if col_mapping is None else col_mapping[col_name]): schema[pxt_name].create_literal(
-                values[i]
-            )
+            (pxt_name := col_mapping.get(col_name, col_name)): _ar_val_to_pxt_val(values[i], schema[pxt_name])
             for col_name, values in pydict.items()
         }

pixeltable/utils/console_output.py

@@ -1,4 +1,5 @@
 import logging
+from typing import TextIO
 
 
 def map_level(verbosity: int) -> int:
@@ -22,10 +23,10 @@ def map_level(verbosity: int) -> int:
 
 
 class ConsoleOutputHandler(logging.StreamHandler):
-    def __init__(self, stream):
+    def __init__(self, stream: TextIO):
         super().__init__(stream)
 
-    def emit(self, record):
+    def emit(self, record: logging.LogRecord) -> None:
         if record.msg.endswith('\n'):
             self.stream.write(record.msg)
         else:

pixeltable/utils/coroutine.py

@@ -7,8 +7,8 @@ T = TypeVar('T')
 
 
 # TODO This is a temporary hack to be able to run async UDFs in contexts that are not properly handled by the existing
-#   scheduler logic (e.g., inside the eval loop of a JsonMapper). Once the scheduler is fully general, it can be
-#   removed.
+#   scheduler logic (e.g., as an embedding function as part of a similarity lookup). Once the scheduler is fully
+#   general, it can be removed.
 
 
 def run_coroutine_synchronously(coroutine: Coroutine[Any, Any, T], timeout: float = 30) -> T:
@@ -16,7 +16,7 @@ def run_coroutine_synchronously(coroutine: Coroutine[Any, Any, T], timeout: floa
     Runs the given coroutine synchronously, even if called in the context of a running event loop.
     """
 
-    def run_in_new_loop():
+    def run_in_new_loop() -> T:
         new_loop = asyncio.new_event_loop()
         asyncio.set_event_loop(new_loop)
         try:

pixeltable/utils/dbms.py

@@ -0,0 +1,66 @@
+import abc
+
+from sqlalchemy import URL
+
+
+class Dbms(abc.ABC):
+    """
+    Provides abstractions for utilities to interact with a database system.
+    """
+
+    name: str
+    transaction_isolation_level: str
+    version_index_type: str
+    db_url: URL
+
+    def __init__(self, name: str, transaction_isolation_level: str, version_index_type: str, db_url: URL) -> None:
+        self.name = name
+        self.transaction_isolation_level = transaction_isolation_level
+        self.version_index_type = version_index_type
+        self.db_url = db_url
+
+    @abc.abstractmethod
+    def drop_db_stmt(self, database: str) -> str: ...
+
+    @abc.abstractmethod
+    def create_db_stmt(self, database: str) -> str: ...
+
+    @abc.abstractmethod
+    def default_system_db_url(self) -> str: ...
+
+
+class PostgresqlDbms(Dbms):
+    """
+    Implements utilities to interact with Postgres database.
+    """
+
+    def __init__(self, db_url: URL):
+        super().__init__('postgresql', 'REPEATABLE READ', 'brin', db_url)
+
+    def drop_db_stmt(self, database: str) -> str:
+        return f'DROP DATABASE {database}'
+
+    def create_db_stmt(self, database: str) -> str:
+        return f"CREATE DATABASE {database} ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
+
+    def default_system_db_url(self) -> str:
+        a = self.db_url.set(database='postgres').render_as_string(hide_password=False)
+        return a
+
+
+class CockroachDbms(Dbms):
+    """
+    Implements utilities to interact with CockroachDb database.
+    """
+
+    def __init__(self, db_url: URL):
+        super().__init__('cockroachdb', 'SERIALIZABLE', 'btree', db_url)
+
+    def drop_db_stmt(self, database: str) -> str:
+        return f'DROP DATABASE {database} CASCADE'
+
+    def create_db_stmt(self, database: str) -> str:
+        return f"CREATE DATABASE {database} TEMPLATE template0 ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C'"
+
+    def default_system_db_url(self) -> str:
+        return self.db_url.set(database='defaultdb').render_as_string(hide_password=False)

pixeltable/utils/documents.py

@@ -1,11 +1,12 @@
 import dataclasses
+import os
 from typing import Optional
 
 import bs4
 import fitz  # type: ignore[import-untyped]
 import puremagic
 
-import pixeltable.type_system as ts
+from pixeltable import exceptions as excs, type_system as ts
 from pixeltable.env import Env
 
 
@@ -18,85 +19,78 @@ class DocumentHandle:
     txt_doc: Optional[str] = None
 
 
-def get_document_handle(path: str) -> Optional[DocumentHandle]:
-    doc_format = puremagic.from_file(path)
+def get_document_handle(path: str) -> DocumentHandle:
+    _, extension = os.path.splitext(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.pdf':
-        pdf_doc = get_pdf_handle(path)
-        if pdf_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.PDF, pdf_doc=pdf_doc)
+    # if no extension, use puremagic to determine the type
+    extension = puremagic.from_file(path)
+    handle = get_handle_by_extension(path, extension)
+    if handle is not None:
+        return handle
 
-    if doc_format == '.html':
-        bs_doc = get_html_handle(path)
-        if bs_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.HTML, bs_doc=bs_doc)
+    raise excs.Error(f'Unrecognized document format: {path}')
 
-    if doc_format == '.md':
-        md_ast = get_markdown_handle(path)
-        if md_ast is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.MD, md_ast=md_ast)
 
-    if doc_format == '.xml':
-        bs_doc = get_xml_handle(path)
-        if bs_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.XML, bs_doc=bs_doc)
+def get_handle_by_extension(path: str, extension: str) -> Optional[DocumentHandle]:
+    doc_format = ts.DocumentType.DocumentFormat.from_extension(extension)
 
-    if doc_format == '.txt':
-        txt_doc = get_txt(path)
-        if txt_doc is not None:
-            return DocumentHandle(format=ts.DocumentType.DocumentFormat.TXT, txt_doc=txt_doc)
+    try:
+        if doc_format == ts.DocumentType.DocumentFormat.HTML:
+            return DocumentHandle(doc_format, bs_doc=get_html_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.MD:
+            return DocumentHandle(doc_format, md_ast=get_markdown_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.PDF:
+            return DocumentHandle(doc_format, pdf_doc=get_pdf_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.XML:
+            return DocumentHandle(doc_format, bs_doc=get_xml_handle(path))
+        if doc_format == ts.DocumentType.DocumentFormat.TXT:
+            return DocumentHandle(doc_format, txt_doc=get_txt(path))
+    except Exception as exc:
+        raise excs.Error(f'An error occurred processing a {doc_format} document: {path}') from exc
 
     return None
 
 
-def get_pdf_handle(path: str) -> Optional[fitz.Document]:
-    try:
-        doc = fitz.open(path)
-        # check pdf (bc it will work for images)
-        if not doc.is_pdf:
-            return None
-        # try to read one page
-        next(page for page in doc)
-        return doc
-    except Exception:
-        return None
-
-
-def get_html_handle(path: str) -> Optional[bs4.BeautifulSoup]:
-    try:
-        with open(path, 'r', encoding='utf8') as fp:
-            doc = bs4.BeautifulSoup(fp, 'lxml')
-        return doc if doc.find() is not None else None
-    except Exception:
-        return None
-
-
-def get_xml_handle(path: str) -> Optional[bs4.BeautifulSoup]:
-    try:
-        with open(path, 'r', encoding='utf8') as fp:
-            doc = bs4.BeautifulSoup(fp, 'xml')
-        return doc if doc.find() is not None else None
-    except Exception:
-        return None
+def get_html_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'lxml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid HTML document: {path}')
+    return doc
 
 
-def get_markdown_handle(path: str) -> Optional[dict]:
+def get_markdown_handle(path: str) -> dict:
     Env.get().require_package('mistune', [3, 0])
     import mistune
 
-    try:
-        with open(path, encoding='utf8') as file:
-            text = file.read()
-        md_ast = mistune.create_markdown(renderer=None)
-        return md_ast(text)
-    except Exception:
-        return None
+    with open(path, encoding='utf8') as file:
+        text = file.read()
+    md_ast = mistune.create_markdown(renderer=None)
+    return md_ast(text)
 
 
-def get_txt(path: str) -> Optional[str]:
-    try:
-        with open(path, 'r', encoding='utf-8') as fp:
-            doc = fp.read()
-        return doc or None  # replace '' with None
-    except Exception:
-        return None
+def get_pdf_handle(path: str) -> fitz.Document:
+    doc = fitz.open(path)
+    # check pdf (bc it will work for images)
+    if not doc.is_pdf:
+        raise excs.Error(f'Not a valid PDF document: {path}')
+    # try to read one page
+    next(page for page in doc)
+    return doc
+
+
+def get_xml_handle(path: str) -> bs4.BeautifulSoup:
+    with open(path, 'r', encoding='utf8') as fp:
+        doc = bs4.BeautifulSoup(fp, 'xml')
+    if doc.find() is None:
+        raise excs.Error(f'Not a valid XML document: {path}')
+    return doc
+
+
+def get_txt(path: str) -> str:
+    with open(path, 'r', encoding='utf-8') as fp:
+        doc = fp.read()
+    return doc

pixeltable/utils/filecache.py

@@ -102,7 +102,7 @@ class FileCache:
     def init(cls) -> None:
         cls.__instance = cls()
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.cache = OrderedDict()
         self.total_size = 0
         self.capacity_bytes = int(Env.get()._file_cache_size_g * (1 << 30))

pixeltable/utils/http_server.py

@@ -3,6 +3,7 @@ import http.server
 import logging
 import pathlib
 import urllib
+from typing import Any
 
 _logger = logging.getLogger('pixeltable.http.server')
 
@@ -38,7 +39,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
         path = pathlib.Path(urllib.request.url2pathname(path))
         return str(path)
 
-    def log_message(self, format, *args) -> None:
+    def log_message(self, format: str, *args: Any) -> None:
         """override logging to stderr in http.server.BaseHTTPRequestHandler"""
         message = format % args
         _logger.info(message.translate(self._control_char_table))  # type: ignore[attr-defined]
@@ -47,7 +48,7 @@ class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
 class LoggingHTTPServer(http.server.ThreadingHTTPServer):
     """Avoids polluting stdout and stderr"""
 
-    def handle_error(self, request, client_address) -> None:
+    def handle_error(self, request, client_address) -> None:  # type: ignore[no-untyped-def]
         """override socketserver.TCPServer.handle_error which prints directly to sys.stderr"""
         import traceback
 

pixeltable/utils/pytorch.py

@@ -32,7 +32,7 @@ class PixeltablePytorchDataset(torch.utils.data.IterableDataset):
 
         self.path = path
         self.image_format = image_format
-        assert image_format in {'np', 'pt'}
+        assert image_format in ('np', 'pt')
         column_type_path = path / '.pixeltable.column_types.json'
         assert column_type_path.exists(), f'missing {column_type_path}'
         with column_type_path.open() as f:

pixeltable/utils/sql.py

@@ -4,7 +4,7 @@ import sqlalchemy as sql
 from sqlalchemy.dialects import postgresql
 
 
-def log_stmt(logger: logging.Logger, stmt) -> None:
+def log_stmt(logger: logging.Logger, stmt: sql.sql.ClauseElement) -> None:
     logger.debug(f'executing {stmt.compile(dialect=postgresql.dialect())}')