pixeltable 0.2.17__py3-none-any.whl → 0.2.18__py3-none-any.whl

pixeltable/io/parquet.py

@@ -7,7 +7,7 @@ import random
 import typing
 from collections import deque
 from pathlib import Path
-from typing import Dict, Optional
+from typing import Dict, Optional, Any
 
 import PIL.Image
 import numpy as np
@@ -142,21 +142,22 @@ def import_parquet(
     table_path: str,
     *,
     parquet_path: str,
-    schema_override: Optional[Dict[str, ts.ColumnType]] = None,
-    **kwargs,
-) -> pxt.catalog.InsertableTable:
-    """Create a new `Table` from a Parquet file or set of files. Requires pyarrow to be installed.
+    schema_overrides: Optional[Dict[str, ts.ColumnType]] = None,
+    **kwargs: Any,
+) -> pxt.Table:
+    """Creates a new base table from a Parquet file or set of files. Requires pyarrow to be installed.
+
     Args:
-        path_str: Path to the table within pixeltable.
+        table_path: Path to the table.
         parquet_path: Path to an individual Parquet file or directory of Parquet files.
-        schema_override: Optional dictionary mapping column names to column type to override the default
-                        schema inferred from the Parquet file. The column type should be a pixeltable ColumnType.
-                        For example, {'col_vid': VideoType()}, rather than {'col_vid': StringType()}.
-                        Any fields not provided explicitly will map to types with `pixeltable.utils.parquet.parquet_schema_to_pixeltable_schema`
+        schema_overrides: If specified, then for each (name, type) pair in `schema_overrides`, the column with
+            name `name` will be given type `type`, instead of being inferred from the Parquet dataset. The keys in
+            `schema_overrides` should be the column names of the Parquet dataset (whether or not they are valid
+            Pixeltable identifiers).
         kwargs: Additional arguments to pass to `create_table`.
 
     Returns:
-        The newly created table. The table will have loaded the data from the Parquet file(s).
+        A handle to the newly created [`Table`][pixeltable.Table].
     """
     import pixeltable as pxt
     from pyarrow import parquet
@@ -166,10 +167,10 @@ def import_parquet(
     parquet_dataset = parquet.ParquetDataset(input_path)
 
     schema = parquet_schema_to_pixeltable_schema(parquet_path)
-    if schema_override is None:
-        schema_override = {}
+    if schema_overrides is None:
+        schema_overrides = {}
 
-    schema.update(schema_override)
+    schema.update(schema_overrides)
     for k, v in schema.items():
         if v is None:
             raise exc.Error(f'Could not infer pixeltable type for column {k} from parquet file')

pixeltable/iterators/document.py

@@ -1,14 +1,15 @@
 import dataclasses
 import enum
 import logging
-from typing import Dict, Any, List, Tuple, Optional, Iterable, Iterator
+from typing import Any, Iterable, Iterator, Optional
 
 import ftfy
 
 from pixeltable.env import Env
 from pixeltable.exceptions import Error
-from pixeltable.type_system import ColumnType, DocumentType, StringType, IntType, JsonType
+from pixeltable.type_system import ColumnType, DocumentType, IntType, JsonType, StringType
 from pixeltable.utils.documents import get_document_handle
+
 from .base import ComponentIterator
 
 _logger = logging.getLogger('pixeltable')
@@ -38,12 +39,12 @@ class DocumentSectionMetadata:
     sourceline: Optional[int] = None
     # the stack of headings up to the most recently observed one;
     # eg, if the most recent one was an h2, 'headings' would contain keys 1 and 2, but nothing below that
-    heading: Optional[Dict[str, str]] = None
+    heading: Optional[dict[str, str]] = None
 
     # pdf-specific metadata
     page: Optional[int] = None
     # bounding box as an {x1, y1, x2, y2} dictionary
-    bounding_box: Optional[Dict[str, float]] = None
+    bounding_box: Optional[dict[str, float]] = None
 
 
 @dataclasses.dataclass
@@ -53,7 +54,7 @@ class DocumentSection:
     metadata: Optional[DocumentSectionMetadata]
 
 
-def _parse_separators(separators: str) -> List[Separator]:
+def _parse_separators(separators: str) -> list[Separator]:
     ret = []
     for s in separators.split(','):
         clean_s = s.strip().upper()
@@ -67,7 +68,7 @@ def _parse_separators(separators: str) -> List[Separator]:
     return ret
 
 
-def _parse_metadata(metadata: str) -> List[ChunkMetadata]:
+def _parse_metadata(metadata: str) -> list[ChunkMetadata]:
     ret = []
     for m in metadata.split(','):
         clean_m = m.strip().upper()
@@ -161,7 +162,7 @@ class DocumentSplitter(ComponentIterator):
             self._sections = self._char_chunks(self._sections)
 
     @classmethod
-    def input_schema(cls) -> Dict[str, ColumnType]:
+    def input_schema(cls) -> dict[str, ColumnType]:
         return {
             'document': DocumentType(nullable=False),
             'separators': StringType(nullable=False),
@@ -174,7 +175,7 @@ class DocumentSplitter(ComponentIterator):
         }
 
     @classmethod
-    def output_schema(cls, *args: Any, **kwargs: Any) -> Tuple[Dict[str, ColumnType], List[str]]:
+    def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ColumnType], list[str]]:
         schema = {'text': StringType()}
         md_fields = _parse_metadata(kwargs['metadata']) if 'metadata' in kwargs else []
 
@@ -208,7 +209,7 @@ class DocumentSplitter(ComponentIterator):
 
         return schema, []
 
-    def __next__(self) -> Dict[str, Any]:
+    def __next__(self) -> dict[str, Any]:
         while True:
             section = next(self._sections)
             if section.text is None:
@@ -236,7 +237,7 @@ class DocumentSplitter(ComponentIterator):
         accumulated_text = []  # currently accumulated text
         # accumulate pieces then join before emit to avoid quadratic complexity of string concatenation
 
-        headings: Dict[str, str] = {}   # current state of observed headings (level -> text)
+        headings: dict[str, str] = {}   # current state of observed headings (level -> text)
         sourceline = 0  # most recently seen sourceline
 
         def update_metadata(el: bs4.Tag) -> None:
@@ -250,7 +251,7 @@ class DocumentSplitter(ComponentIterator):
                     del headings[l]
                 headings[el.name] = el.get_text().strip()
 
-        def emit() -> None:
+        def emit() -> Iterator[DocumentSection]:
             nonlocal accumulated_text, headings, sourceline
             if len(accumulated_text) > 0:
                 md = DocumentSectionMetadata(sourceline=sourceline, heading=headings.copy())
@@ -294,9 +295,9 @@ class DocumentSplitter(ComponentIterator):
         # current state
         accumulated_text = []  # currently accumulated text
         # accumulate pieces then join before emit to avoid quadratic complexity of string concatenation
-        headings: Dict[str, str] = {}   # current state of observed headings (level -> text)
+        headings: dict[str, str] = {}   # current state of observed headings (level -> text)
 
-        def update_headings(heading: Dict) -> None:
+        def update_headings(heading: dict) -> None:
             # update current state
             nonlocal headings
             assert 'type' in heading and heading['type'] == 'heading'
@@ -309,14 +310,14 @@ class DocumentSplitter(ComponentIterator):
                 del headings[l]
             headings[level] = text
 
-        def emit() -> None:
+        def emit() -> Iterator[DocumentSection]:
             nonlocal accumulated_text, headings
             if len(accumulated_text) > 0:
                 metadata = DocumentSectionMetadata(sourceline=0, heading=headings.copy())
                 yield DocumentSection(text=ftfy.fix_text(' '.join(accumulated_text)), metadata=metadata)
                 accumulated_text = []
 
-        def process_element(el: Dict) -> Iterator[DocumentSection]:
+        def process_element(el: dict) -> Iterator[DocumentSection]:
             # process the element and emit sections as necessary
             nonlocal accumulated_text, headings, emit_on_heading, emit_on_paragraph
             assert 'type' in el

pixeltable/metadata/__init__.py

@@ -10,7 +10,7 @@ import sqlalchemy.orm as orm
 from .schema import SystemInfo, SystemInfoMd
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 19
+VERSION = 20
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:

pixeltable/metadata/converters/convert_19.py

@@ -0,0 +1,46 @@
+import datetime
+from typing import Any, Optional
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.metadata import register_converter, schema
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=19)
+def _(engine: sql.engine.Engine) -> None:
+    # Convert all timestamp literals to aware datetimes
+    convert_table_md(engine, substitution_fn=__update_timestamp_literals)
+
+    # Convert all timestamp columns to TIMESTAMPTZ. (This conversion will take place in the database
+    # default time zone, which is what we want, since in versions <= 19 they were naive timestamps.)
+    with engine.begin() as conn:
+        tables = conn.execute(sql.select(schema.Table.id, schema.Table.md))
+        for id, md in tables:
+            store_prefix = 'view' if md['view_md'] is not None else 'tbl'
+            store_name = f'{store_prefix}_{id.hex}'
+            column_md = md['column_md']
+            timestamp_cols = [
+                col_id for col_id, col in column_md.items()
+                if col['col_type']['_classname'] == 'TimestampType'
+            ]
+            for col_id in timestamp_cols:
+                conn.execute(
+                    sql.text(f'ALTER TABLE {store_name} ALTER COLUMN col_{col_id} TYPE TIMESTAMPTZ')
+                )
+
+
+def __update_timestamp_literals(k: Any, v: Any) -> Optional[tuple[Any, Any]]:
+    if isinstance(v, dict) and 'val_t' in v:
+        # It's a literal with an explicit 'val_t' field. In version 19 this can only mean a
+        # timestamp literal, which (in version 19) is stored in the DB as a naive datetime.
+        # We convert it to an aware datetime, stored in UTC.
+        assert v['_classname'] == 'Literal'
+        assert v['val_t'] == pxt.ColumnType.Type.TIMESTAMP.name
+        assert isinstance(v['val'], str)
+        dt = datetime.datetime.fromisoformat(v['val'])
+        assert dt.tzinfo is None  # In version 19 all timestamps are naive
+        dt_utc = dt.astimezone(datetime.timezone.utc)
+        v['val'] = dt_utc.isoformat()
+        return k, v

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 = {
+    20: 'Store DB timestamps in UTC',
     19: 'UDF renames; ImageMemberAccess removal',
     18: 'Restructured index metadata',
     17: 'Renamed remotes to external_stores',

pixeltable/metadata/schema.py

@@ -3,6 +3,7 @@ import uuid
 from typing import Optional, List, get_type_hints, Type, Any, TypeVar, Tuple, Union
 
 import sqlalchemy as sql
+import sqlalchemy.orm as orm
 from sqlalchemy import ForeignKey
 from sqlalchemy import Integer, BigInteger, LargeBinary
 from sqlalchemy.dialects.postgresql import UUID, JSONB
@@ -64,8 +65,8 @@ class DirMd:
 class Dir(Base):
     __tablename__ = 'dirs'
 
-    id = sql.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
-    parent_id = sql.Column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
+    id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, nullable=False)
+    parent_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=True)
     md = sql.Column(JSONB, nullable=False)
 
 
@@ -163,8 +164,8 @@ class Table(Base):
 
     MAX_VERSION = 9223372036854775807  # 2^63 - 1
 
-    id = sql.Column(UUID(as_uuid=True), primary_key=True, nullable=False)
-    dir_id = sql.Column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=False)
+    id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), primary_key=True, nullable=False)
+    dir_id: orm.Mapped[uuid.UUID] = orm.mapped_column(UUID(as_uuid=True), ForeignKey('dirs.id'), nullable=False)
     md = sql.Column(JSONB, nullable=False)  # TableMd