pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exprs/json_path.py

@@ -1,59 +1,75 @@
 from __future__ import annotations
 
-from typing import Any, Optional, Union
+import io
+from pathlib import Path
+from typing import Any
 
 import jmespath
 import sqlalchemy as sql
 
-import pixeltable as pxt
-import pixeltable.catalog as catalog
-import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import catalog, exceptions as excs, type_system as ts
 
+from .column_ref import ColumnRef
 from .data_row import DataRow
 from .expr import Expr
 from .globals import print_slice
-from .json_mapper import JsonMapper
+from .json_mapper import JsonMapperDispatch
+from .object_ref import ObjectRef
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
 
 class JsonPath(Expr):
+    """
+    anchor can be None, in which case this is a relative JsonPath and the anchor is set later via set_anchor().
+    scope_idx: for relative paths, index of referenced JsonMapper
+    (0: indicates the immediately preceding JsonMapper, -1: the parent of the immediately preceding mapper, ...)
+    """
+
+    path_elements: list[str | int | slice]
+    compiled_path: jmespath.parser.ParsedResult | None
+    scope_idx: int
+    file_handles: dict[Path, io.BufferedReader]  # key: file path
+
     def __init__(
-        self,
-        anchor: Optional['pxt.exprs.Expr'],
-        path_elements: Optional[list[Union[str, int, slice]]] = None,
-        scope_idx: int = 0
+        self, anchor: Expr | None, path_elements: list[str | int | slice] | None = None, scope_idx: int = 0
     ) -> None:
-        """
-        anchor can be None, in which case this is a relative JsonPath and the anchor is set later via set_anchor().
-        scope_idx: for relative paths, index of referenced JsonMapper
-        (0: indicates the immediately preceding JsonMapper, -1: the parent of the immediately preceding mapper, ...)
-        """
         if path_elements is None:
             path_elements = []
         super().__init__(ts.JsonType(nullable=True))  # JsonPath expressions are always nullable
         if anchor is not None:
             self.components = [anchor]
-        self.path_elements: list[Union[str, int, slice]] = path_elements
+        self.path_elements = path_elements
         self.compiled_path = jmespath.compile(self._json_path()) if len(path_elements) > 0 else None
         self.scope_idx = scope_idx
         # NOTE: the _create_id() result will change if set_anchor() gets called;
         # this is not a problem, because _create_id() shouldn't be called after init()
         self.id = self._create_id()
+        self.file_handles = {}
+
+    def release(self) -> None:
+        for fh in self.file_handles.values():
+            fh.close()
+        self.file_handles.clear()
 
     def __repr__(self) -> str:
-        # else "R": the anchor is RELATIVE_PATH_ROOT
-        return (f'{str(self._anchor) if self._anchor is not None else "R"}'
-            f'{"." if isinstance(self.path_elements[0], str) else ""}{self._json_path()}')
+        # else 'R': the anchor is RELATIVE_PATH_ROOT
+        anchor_str = str(self.anchor) if self.anchor is not None else 'R'
+        if len(self.path_elements) == 0:
+            return anchor_str
+        return f'{anchor_str}{"." if isinstance(self.path_elements[0], str) else ""}{self._json_path()}'
 
     def _as_dict(self) -> dict:
-        path_elements = [
-            [el.start, el.stop, el.step] if isinstance(el, slice)
-            else el
-            for el in self.path_elements
-        ]
-        return {'path_elements': path_elements, 'scope_idx': self.scope_idx, **super()._as_dict()}
+        assert len(self.components) <= 1
+        components_dict: dict[str, Any]
+        if len(self.components) == 0 or isinstance(self.components[0], ObjectRef):
+            # If the anchor is an ObjectRef, it means this JsonPath is a bound relative path. We store it as a relative
+            # path, *not* a bound path (which has no meaning in the dict).
+            components_dict = {}
+        else:
+            components_dict = super()._as_dict()
+        path_elements = [[el.start, el.stop, el.step] if isinstance(el, slice) else el for el in self.path_elements]
+        return {'path_elements': path_elements, 'scope_idx': self.scope_idx, **components_dict}
 
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> JsonPath:
@@ -61,15 +77,11 @@ class JsonPath(Expr):
         assert 'scope_idx' in d
         assert len(components) <= 1
         anchor = components[0] if len(components) == 1 else None
-        path_elements = [
-            slice(el[0], el[1], el[2]) if isinstance(el, list)
-            else el
-            for el in d['path_elements']
-        ]
+        path_elements = [slice(el[0], el[1], el[2]) if isinstance(el, list) else el for el in d['path_elements']]
         return cls(anchor, path_elements, d['scope_idx'])
 
     @property
-    def _anchor(self) -> Optional[Expr]:
+    def anchor(self) -> Expr | None:
         return None if len(self.components) == 0 else self.components[0]
 
     def set_anchor(self, anchor: Expr) -> None:
@@ -77,44 +89,42 @@ class JsonPath(Expr):
         self.components = [anchor]
 
     def is_relative_path(self) -> bool:
-        return self._anchor is None
+        return self.anchor is None
 
-    def bind_rel_paths(self, mapper: Optional['JsonMapper'] = None) -> None:
-        if not self.is_relative_path():
-            return
-        # TODO: take scope_idx into account
-        self.set_anchor(mapper.scope_anchor)
+    def _has_relative_path(self) -> bool:
+        return self.is_relative_path() or super()._has_relative_path()
+
+    def _bind_rel_paths(self, mapper: 'JsonMapperDispatch' | None = None) -> None:
+        if self.is_relative_path():
+            # TODO: take scope_idx into account
+            self.set_anchor(mapper.scope_anchor)
+        else:
+            self.anchor._bind_rel_paths(mapper)
 
     def __call__(self, *args: object, **kwargs: object) -> 'JsonPath':
         """
         Construct a relative path that references an ancestor of the immediately enclosing JsonMapper.
         """
         if not self.is_relative_path():
-            raise excs.Error(f'() for an absolute path is invalid')
+            raise excs.Error('() for an absolute path is invalid')
         if len(args) != 1 or not isinstance(args[0], int) or args[0] >= 0:
-            raise excs.Error(f'R() requires a negative index')
+            raise excs.Error('R() requires a negative index')
         return JsonPath(None, [], args[0])
 
     def __getattr__(self, name: str) -> 'JsonPath':
         assert isinstance(name, str)
-        return JsonPath(self._anchor, self.path_elements + [name])
+        return JsonPath(self.anchor, [*self.path_elements, name])
 
     def __getitem__(self, index: object) -> 'JsonPath':
         if isinstance(index, (int, slice, str)):
-            return JsonPath(self._anchor, self.path_elements + [index])
+            return JsonPath(self.anchor, [*self.path_elements, index])
         raise excs.Error(f'Invalid json list index: {index}')
 
-    def __rshift__(self, other: object) -> 'JsonMapper':
-        rhs_expr = Expr.from_object(other)
-        if rhs_expr is None:
-            raise excs.Error(f'>> requires an expression on the right-hand side, found {type(other)}')
-        return JsonMapper(self, rhs_expr)
-
-    def default_column_name(self) -> Optional[str]:
-        anchor_name = self._anchor.default_column_name() if self._anchor is not None else ''
+    def default_column_name(self) -> str | None:
+        anchor_name = self.anchor.default_column_name() if self.anchor is not None else ''
         ret_name = f'{anchor_name}.{self._json_path()}'
 
-        def cleanup_char(s : str) -> str:
+        def cleanup_char(s: str) -> str:
             if s == '.':
                 return '_'
             elif s == '*':
@@ -125,8 +135,8 @@ class JsonPath(Expr):
                 return ''
 
         clean_name = ''.join(map(cleanup_char, ret_name))
-        clean_name = clean_name.lstrip('_') # remove leading underscore
-        if clean_name == '':
+        clean_name = clean_name.lstrip('_')  # remove leading underscore
+        if not clean_name:  # Replace '' with None
             clean_name = None
 
         assert clean_name is None or catalog.is_valid_identifier(clean_name)
@@ -136,17 +146,17 @@ class JsonPath(Expr):
         return self.path_elements == other.path_elements
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('path_elements', self.path_elements)]
+        return [*super()._id_attrs(), ('path_elements', self.path_elements)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         """
         Postgres appears to have a bug: jsonb_path_query('{a: [{b: 0}, {b: 1}]}', '$.a.b') returns
         *two* rows (each containing col val 0), not a single row with [0, 0].
         We need to use a workaround: retrieve the entire dict, then use jmespath to extract the path correctly.
         """
-        #path_str = '$.' + '.'.join(self.path_elements)
-        #assert isinstance(self._anchor(), ColumnRef)
-        #return sql.func.jsonb_path_query(self._anchor().col.sa_col, path_str)
+        # path_str = '$.' + '.'.join(self.path_elements)
+        # assert isinstance(self._anchor(), ColumnRef)
+        # return sql.func.jsonb_path_query(self._anchor().col.sa_col, path_str)
         return None
 
     def _json_path(self) -> str:
@@ -163,11 +173,31 @@ class JsonPath(Expr):
                 result.append(f'[{print_slice(element)}]')
         return ''.join(result)
 
-    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        val = data_row[self._anchor.slot_idx]
+    def eval(self, row: DataRow, row_builder: RowBuilder) -> None:
+        assert self.anchor is not None, self
+        val = row[self.anchor.slot_idx]
         if self.compiled_path is not None:
             val = self.compiled_path.search(val)
-        data_row[self.slot_idx] = val
+        row[self.slot_idx] = val
+        if val is None or self.anchor is None or not isinstance(self.anchor, ColumnRef):
+            return
+
+        # the origin of val is a json-typed column, which might stored inlined objects
+        if self.anchor.slot_idx not in row.slot_md:
+            # we can infer that there aren't any inlined objects because our execution plan doesn't include
+            # materializing the cellmd (eg, insert plans)
+            # TODO: have the planner pass that fact into ExprEvalNode explicitly to streamline this path a bit more
+            return
+
+        # defer import until it's needed
+        from pixeltable.exec.cell_reconstruction_node import json_has_inlined_objs, reconstruct_json
+
+        cell_md = row.slot_md[self.anchor.slot_idx]
+        if cell_md is None or cell_md.file_urls is None or not json_has_inlined_objs(val):
+            # val doesn't contain inlined objects
+            return
+
+        row.vals[self.slot_idx] = reconstruct_json(val, cell_md.file_urls, self.file_handles)
 
 
 RELATIVE_PATH_ROOT = JsonPath(None)

pixeltable/exprs/literal.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+import base64
 import datetime
-from typing import Any, Optional
+import uuid
+from typing import Any
 
+import numpy as np
 import sqlalchemy as sql
 
 import pixeltable.type_system as ts
@@ -15,7 +18,9 @@ from .sql_element_cache import SqlElementCache
 
 
 class Literal(Expr):
-    def __init__(self, val: Any, col_type: Optional[ts.ColumnType] = None):
+    val: Any
+
+    def __init__(self, val: Any, col_type: ts.ColumnType | None = None):
         if col_type is not None:
             val = col_type.create_literal(val)
         else:
@@ -33,10 +38,13 @@ class Literal(Expr):
                     val = val.replace(tzinfo=default_tz)
             # Now convert to UTC
             val = val.astimezone(datetime.timezone.utc)
+        if isinstance(val, tuple):
+            # Tuples are stored as a list
+            val = list(val)
         self.val = val
         self.id = self._create_id()
 
-    def default_column_name(self) -> Optional[str]:
+    def default_column_name(self) -> str | None:
         return 'Literal'
 
     def __str__(self) -> str:
@@ -46,6 +54,15 @@ class Literal(Expr):
             assert isinstance(self.val, datetime.datetime)
             default_tz = Env.get().default_time_zone
             return f"'{self.val.astimezone(default_tz).isoformat()}'"
+        if self.col_type.is_date_type():
+            assert isinstance(self.val, datetime.date)
+            return f"'{self.val.isoformat()}'"
+        if self.col_type.is_uuid_type():
+            assert isinstance(self.val, uuid.UUID)
+            return f"'{self.val}'"
+        if self.col_type.is_array_type():
+            assert isinstance(self.val, np.ndarray)
+            return str(self.val.tolist())
         return str(self.val)
 
     def __repr__(self) -> str:
@@ -55,19 +72,18 @@ class Literal(Expr):
         return self.val == other.val
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('val', self.val)]
+        return [*super()._id_attrs(), ('val', self.val)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
-        # we need to return something here so that we can generate a Where clause for predicates
-        # that involve literals (like Where c > 0)
-        return sql.sql.expression.literal(self.val)
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
+        # Return a sql object so that constants can participate in SQL expressions
+        return sql.sql.expression.literal(self.val, type_=self.col_type.to_sa_type())
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         # this will be called, even though sql_expr() does not return None
         data_row[self.slot_idx] = self.val
 
     def _as_dict(self) -> dict:
-        # For some types, we need to explictly record their type, because JSON does not know
+        # For some types, we need to explicitly record their type, because JSON does not know
         # how to interpret them unambiguously
         if self.col_type.is_timestamp_type():
             assert isinstance(self.val, datetime.datetime)
@@ -76,18 +92,47 @@ class Literal(Expr):
             # stored as UTC in the database)
             encoded_val = self.val.isoformat()
             return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
+        elif self.col_type.is_date_type():
+            assert isinstance(self.val, datetime.date)
+            encoded_val = self.val.isoformat()
+            return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
+        elif self.col_type.is_uuid_type():
+            assert isinstance(self.val, uuid.UUID)
+            encoded_val = str(self.val)
+            return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
+        elif self.col_type.is_binary_type():
+            assert isinstance(self.val, bytes)
+            encoded_val = base64.b64encode(self.val).decode('utf-8')
+            return {'val': encoded_val, 'val_t': self.col_type._type.name, **super()._as_dict()}
+        elif self.col_type.is_array_type():
+            assert isinstance(self.val, np.ndarray)
+            return {'val': self.val.tolist(), 'val_t': self.col_type._type.name, **super()._as_dict()}
         else:
             return {'val': self.val, **super()._as_dict()}
 
+    def as_literal(self) -> Literal | None:
+        return self
+
     @classmethod
     def _from_dict(cls, d: dict, components: list[Expr]) -> Literal:
         assert 'val' in d
         if 'val_t' in d:
             val_t = d['val_t']
-            # Currently the only special-cased literal type is TIMESTAMP
-            assert val_t == ts.ColumnType.Type.TIMESTAMP.name
-            dt = datetime.datetime.fromisoformat(d['val'])
-            assert dt.tzinfo == datetime.timezone.utc  # Must be UTC in the database
-            return cls(dt)
-        else:
-            return cls(d['val'])
+            if val_t == ts.ColumnType.Type.DATE.name:
+                dt = datetime.date.fromisoformat(d['val'])
+                return cls(dt)
+            elif val_t == ts.ColumnType.Type.TIMESTAMP.name:
+                dt = datetime.datetime.fromisoformat(d['val'])
+                assert dt.tzinfo == datetime.timezone.utc  # Must be UTC in the database
+                return cls(dt)
+            elif val_t == ts.ColumnType.Type.UUID.name:
+                uuid_val = uuid.UUID(d['val'])
+                return cls(uuid_val)
+            elif val_t == ts.ColumnType.Type.BINARY.name:
+                assert isinstance(d['val'], str)
+                bytes_val = base64.b64decode(d['val'].encode('utf-8'))
+                return cls(bytes_val)
+            elif val_t == ts.ColumnType.Type.ARRAY.name:
+                arrays = np.array(d['val'])
+                return cls(arrays)
+        return cls(d['val'])

pixeltable/exprs/method_ref.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -19,10 +19,11 @@ class MethodRef(Expr):
     When a `MethodRef` is called, it returns a `FunctionCall` with the base expression as the first argument.
     The effective arity of a `MethodRef` is one less than the arity of the underlying `Function`.
     """
+
     # TODO: Should this even be an `Expr`? It can't actually be evaluated directly (it has to be first
     #   converted to a `FunctionCall` by binding any remaining parameters).
 
-    def __init__(self, base_expr: Expr, method_name: str):
+    def __init__(self, base_expr: Expr, method_name: str) -> None:
         super().__init__(ts.InvalidType())  # The `MethodRef` is untyped until it is called.
         self.base_expr = base_expr
         self.method_name = method_name
@@ -42,7 +43,7 @@ class MethodRef(Expr):
         assert len(components) == 1
         return cls(components[0], d['method_name'])
 
-    def __call__(self, *args, **kwargs) -> FunctionCall:
+    def __call__(self, *args: Any, **kwargs: Any) -> FunctionCall:
         result = self.fn(*[self.base_expr, *args], **kwargs)
         assert isinstance(result, FunctionCall)
         result.is_method_call = True
@@ -52,13 +53,13 @@ class MethodRef(Expr):
         return self.base_expr.id == other.base_expr.id and self.method_name == other.method_name
 
     def _id_attrs(self) -> list[tuple[str, Any]]:
-        return super()._id_attrs() + [('method_name', self.method_name)]
+        return [*super()._id_attrs(), ('method_name', self.method_name)]
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
-        assert False, 'MethodRef cannot be evaluated directly'
+        raise AssertionError('MethodRef cannot be evaluated directly')
 
     def __repr__(self) -> str:
         return f'{self.base_expr}.{self.method_name}'

pixeltable/exprs/object_ref.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Any
 
 import sqlalchemy as sql
 
@@ -8,7 +8,7 @@ import pixeltable.type_system as ts
 
 from .data_row import DataRow
 from .expr import Expr, ExprScope
-from .json_mapper import JsonMapper
+from .json_mapper import JsonMapperDispatch
 from .row_builder import RowBuilder
 from .sql_element_cache import SqlElementCache
 
@@ -18,26 +18,37 @@ class ObjectRef(Expr):
     Reference to an intermediate result, such as the "scope variable" produced by a JsonMapper.
     The object is generated/materialized elsewhere and establishes a new scope.
     """
-    def __init__(self, scope: ExprScope, owner: JsonMapper):
+
+    def __init__(self, scope: ExprScope, owner: JsonMapperDispatch):
         # TODO: do we need an Unknown type after all?
         super().__init__(ts.JsonType())  # JsonType: this could be anything
         self._scope = scope
         self.owner = owner
         self.id = self._create_id()
 
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        # We have no components, so we can't rely on the default behavior here (otherwise, all ObjectRef
+        # instances will be conflated into a single slot).
+        return [('addr', id(self))]
+
+    def substitute(self, subs: dict[Expr, Expr]) -> Expr:
+        # Just return self; we need to avoid creating a new id after doing the substitution, because otherwise
+        # we'll wind up in a situation where the scope_anchor of the enclosing JsonMapper is different from the
+        # nested ObjectRefs inside its target_expr (and therefore occupies a different slot_idx).
+        return self
+
     def scope(self) -> ExprScope:
         return self._scope
 
-    def __str__(self) -> str:
-        assert False
-
     def _equals(self, other: ObjectRef) -> bool:
-        return self.owner is other.owner
+        return self.id == other.id
 
-    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+    def sql_expr(self, _: SqlElementCache) -> sql.ColumnElement | None:
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
         # this will be called, but the value has already been materialized elsewhere
         pass
 
+    def __repr__(self) -> str:
+        return f'ObjectRef({self.owner}, {self.id}, {self.owner.id})'