pixeltable 0.3.12__py3-none-any.whl → 0.3.14__py3-none-any.whl

pixeltable/exec/sql_node.py

@@ -424,6 +424,7 @@ class SqlAggregationNode(SqlNode):
     """
 
     group_by_items: Optional[list[exprs.Expr]]
+    input_cte: Optional[sql.CTE]
 
     def __init__(
         self,
@@ -440,13 +441,13 @@ class SqlAggregationNode(SqlNode):
             group_by_items: list of expressions to group by
             limit: max number of rows to return: None = no limit
         """
-        _, input_col_map = input.to_cte()
+        self.input_cte, input_col_map = input.to_cte()
         sql_elements = exprs.SqlElementCache(input_col_map)
         super().__init__(None, row_builder, select_list, sql_elements)
         self.group_by_items = group_by_items
 
     def _create_stmt(self) -> sql.Select:
-        stmt = super()._create_stmt()
+        stmt = super()._create_stmt().select_from(self.input_cte)
         if self.group_by_items is not None:
             sql_group_by_items = [self.sql_elements.get(e) for e in self.group_by_items]
             assert all(e is not None for e in sql_group_by_items)

pixeltable/exprs/column_ref.py

@@ -31,12 +31,18 @@ class ColumnRef(Expr):
     - in that case, the ColumnRef also instantiates a second non-validating ColumnRef as a component (= dependency)
     - the non-validating ColumnRef is used for SQL translation
 
+    A ColumnRef may have an optional reference table, which carries the context of the ColumnRef resolution. Thus
+    if `v` is a view of `t` (for example), then `v.my_col` and `t.my_col` refer to the same underlying column, but
+    their reference tables will be `v` and `t`, respectively. This is to ensure correct behavior of expressions such
+    as `v.my_col.head()`.
+
     TODO:
     separate Exprs (like validating ColumnRefs) from the logical expression tree and instead have RowBuilder
     insert them into the EvalCtxs as needed
     """
 
     col: catalog.Column
+    reference_tbl: Optional[catalog.TableVersionPath]
     is_unstored_iter_col: bool
     iter_arg_ctx: Optional[RowBuilder.EvalCtx]
     base_rowid_len: int
@@ -46,10 +52,16 @@ class ColumnRef(Expr):
     id: int
     perform_validation: bool  # if True, performs media validation
 
-    def __init__(self, col: catalog.Column, perform_validation: Optional[bool] = None):
+    def __init__(
+        self,
+        col: catalog.Column,
+        reference_tbl: Optional[catalog.TableVersionPath] = None,
+        perform_validation: Optional[bool] = None,
+    ):
         super().__init__(col.col_type)
         assert col.tbl is not None
         self.col = col
+        self.reference_tbl = reference_tbl
         self.is_unstored_iter_col = (
             col.tbl.get().is_component_view and col.tbl.get().is_iterator_column(col) and not col.is_stored
         )
@@ -95,7 +107,7 @@ class ColumnRef(Expr):
         target = tbl_versions[self.col.tbl.id]
         assert self.col.id in target.cols_by_id
         col = target.cols_by_id[self.col.id]
-        return ColumnRef(col)
+        return ColumnRef(col, self.reference_tbl)
 
     def __getattr__(self, name: str) -> Expr:
         from .column_property_ref import ColumnPropertyRef
@@ -126,26 +138,26 @@ class ColumnRef(Expr):
 
         return super().__getattr__(name)
 
-    @classmethod
     def find_embedding_index(
-        cls, col: catalog.Column, idx_name: Optional[str], method_name: str
+        self, idx_name: Optional[str], method_name: str
     ) -> dict[str, catalog.TableVersion.IndexInfo]:
         """Return IndexInfo for a column, with an optional given name"""
-        # determine index to use
-        idx_info_dict = col.get_idx_info()
         from pixeltable import index
 
+        # determine index to use
+        idx_info_dict = self.col.get_idx_info(self.reference_tbl)
+
         embedding_idx_info = {
             info: value for info, value in idx_info_dict.items() if isinstance(value.idx, index.EmbeddingIndex)
         }
         if len(embedding_idx_info) == 0:
-            raise excs.Error(f'No indices found for {method_name!r} on column {col.name!r}')
+            raise excs.Error(f'No indices found for {method_name!r} on column {self.col.name!r}')
         if idx_name is not None and idx_name not in embedding_idx_info:
-            raise excs.Error(f'Index {idx_name!r} not found for {method_name!r} on column {col.name!r}')
+            raise excs.Error(f'Index {idx_name!r} not found for {method_name!r} on column {self.col.name!r}')
         if len(embedding_idx_info) > 1:
             if idx_name is None:
                 raise excs.Error(
-                    f'Column {col.name!r} has multiple indices; use the index name to disambiguate: '
+                    f'Column {self.col.name!r} has multiple indices; use the index name to disambiguate: '
                     f'`{method_name}(..., idx=<index_name>)`'
                 )
             idx_info = {idx_name: embedding_idx_info[idx_name]}
@@ -159,7 +171,7 @@ class ColumnRef(Expr):
         return SimilarityExpr(self, item, idx_name=idx)
 
     def embedding(self, *, idx: Optional[str] = None) -> ColumnRef:
-        idx_info = ColumnRef.find_embedding_index(self.col, idx, 'embedding')
+        idx_info = self.find_embedding_index(idx, 'embedding')
         assert len(idx_info) == 1
         col = copy.copy(next(iter(idx_info.values())).val_col)
         col.name = f'{self.col.name}_embedding_{idx if idx is not None else ""}'
@@ -167,14 +179,21 @@ class ColumnRef(Expr):
         return ColumnRef(col)
 
     def default_column_name(self) -> Optional[str]:
-        return str(self)
+        return self.col.name if self.col is not None else None
 
     def _equals(self, other: ColumnRef) -> bool:
         return self.col == other.col and self.perform_validation == other.perform_validation
 
     def _df(self) -> 'pxt.dataframe.DataFrame':
-        tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
-        return tbl.select(self)
+        from pixeltable import plan
+
+        if self.reference_tbl is None:
+            # No reference table; use the current version of the table to which the column belongs
+            tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
+            return tbl.select(self)
+        else:
+            # Explicit reference table; construct a DataFrame directly from it
+            return pxt.DataFrame(plan.FromClause([self.reference_tbl])).select(self)
 
     def show(self, *args: Any, **kwargs: Any) -> 'pxt.dataframe.DataFrameResultSet':
         return self._df().show(*args, **kwargs)
@@ -188,6 +207,10 @@ class ColumnRef(Expr):
     def count(self) -> int:
         return self._df().count()
 
+    def distinct(self) -> 'pxt.dataframe.DataFrame':
+        """Return distinct values in this column."""
+        return self._df().distinct()
+
     def __str__(self) -> str:
         if self.col.name is None:
             return f'<unnamed column {self.col.id}>'
@@ -203,7 +226,7 @@ class ColumnRef(Expr):
     def _descriptors(self) -> DescriptionHelper:
         tbl = catalog.Catalog.get().get_table_by_id(self.col.tbl.id)
         helper = DescriptionHelper()
-        helper.append(f'Column\n{self.col.name!r}\n(of table {tbl._path()!r})')
+        helper.append(f'Column\n{self.col.name!r}\n(of table {tbl._path!r})')
         helper.append(tbl._col_descriptor([self.col.name]))
         idxs = tbl._index_descriptor([self.col.name])
         if len(idxs) > 0:
@@ -260,13 +283,14 @@ class ColumnRef(Expr):
 
     def _as_dict(self) -> dict:
         tbl = self.col.tbl
-        version = tbl.get().version if tbl.get().is_snapshot else None
+        tbl_version = tbl.get().version if tbl.get().is_snapshot else None
         # we omit self.components, even if this is a validating ColumnRef, because init() will recreate the
         # non-validating component ColumnRef
         return {
             'tbl_id': str(tbl.id),
-            'tbl_version': version,
+            'tbl_version': tbl_version,
             'col_id': self.col.id,
+            'reference_tbl': self.reference_tbl.as_dict() if self.reference_tbl is not None else None,
             'perform_validation': self.perform_validation,
         }
 
@@ -281,5 +305,6 @@ class ColumnRef(Expr):
     @classmethod
     def _from_dict(cls, d: dict, _: list[Expr]) -> ColumnRef:
         col = cls.get_column(d)
+        reference_tbl = None if d['reference_tbl'] is None else catalog.TableVersionPath.from_dict(d['reference_tbl'])
         perform_validation = d['perform_validation']
-        return cls(col, perform_validation=perform_validation)
+        return cls(col, reference_tbl, perform_validation=perform_validation)

pixeltable/exprs/data_row.py

@@ -29,10 +29,13 @@ class DataRow:
     - FloatType: float
     - BoolType: bool
     - TimestampType: datetime.datetime
+    - DateType: datetime.date
     - JsonType: json-serializable object
     - ArrayType: numpy.ndarray
     - ImageType: PIL.Image.Image
     - VideoType: local path if available, otherwise url
+    - AudioType: local path if available, otherwise url
+    - DocumentType: local path if available, otherwise url
     """
 
     vals: np.ndarray  # of object

pixeltable/exprs/globals.py

@@ -5,7 +5,7 @@ import enum
 from typing import Union
 
 # Python types corresponding to our literal types
-LiteralPythonTypes = Union[str, int, float, bool, datetime.datetime]
+LiteralPythonTypes = Union[str, int, float, bool, datetime.datetime, datetime.date]
 
 
 def print_slice(s: slice) -> str:

pixeltable/exprs/literal.py

@@ -50,6 +50,9 @@ 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_array_type():
             assert isinstance(self.val, np.ndarray)
             return str(self.val.tolist())
@@ -82,6 +85,10 @@ 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_array_type():
             assert isinstance(self.val, np.ndarray)
             return {'val': self.val.tolist(), 'val_t': self.col_type._type.name, **super()._as_dict()}
@@ -96,7 +103,10 @@ class Literal(Expr):
         assert 'val' in d
         if 'val_t' in d:
             val_t = d['val_t']
-            if val_t == ts.ColumnType.Type.TIMESTAMP.name:
+            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)

pixeltable/exprs/rowid_ref.py

@@ -30,7 +30,7 @@ class RowidRef(Expr):
 
     def __init__(
         self,
-        tbl: catalog.TableVersionHandle,
+        tbl: Optional[catalog.TableVersionHandle],
         idx: int,
         tbl_id: Optional[UUID] = None,
         normalized_base_id: Optional[UUID] = None,
@@ -98,6 +98,9 @@ class RowidRef(Expr):
     def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
         tbl = self.tbl.get() if self.tbl is not None else catalog.Catalog.get().get_tbl_version(self.tbl_id, None)
         rowid_cols = tbl.store_tbl.rowid_columns()
+        assert self.rowid_component_idx <= len(rowid_cols), (
+            f'{self.rowid_component_idx} not consistent with {rowid_cols}'
+        )
         return rowid_cols[self.rowid_component_idx]
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/exprs/similarity_expr.py

@@ -26,7 +26,7 @@ class SimilarityExpr(Expr):
         from pixeltable import index
 
         # determine index to use
-        idx_dict = ColumnRef.find_embedding_index(col_ref.col, idx_name, 'similarity')
+        idx_dict = col_ref.find_embedding_index(idx_name, 'similarity')
         assert len(idx_dict) == 1
         self.idx_info = next(iter(idx_dict.values()))
         idx = self.idx_info.idx

pixeltable/func/function.py

@@ -514,7 +514,7 @@ class InvalidFunction(Function):
     def _as_dict(self) -> dict:
         """
         Here we write out (verbatim) the original metadata that failed to load (and that resulted in the
-        InvalidFunction). Note that the InvalidFunction itself is never serlialized, so there is no corresponding
+        InvalidFunction). Note that the InvalidFunction itself is never serialized, so there is no corresponding
         from_dict() method.
         """
         return self.fn_dict

pixeltable/func/udf.py

@@ -262,7 +262,7 @@ def from_table(
     """
     from pixeltable import exprs
 
-    ancestors = [tbl, *tbl._bases]
+    ancestors = [tbl, *tbl._base_tables]
     ancestors.reverse()  # We must traverse the ancestors in order from base to derived
 
     subst: dict[exprs.Expr, exprs.Expr] = {}

pixeltable/functions/__init__.py

@@ -5,6 +5,8 @@ from pixeltable.utils.code import local_public_names
 from . import (
     anthropic,
     audio,
+    bedrock,
+    date,
     deepseek,
     fireworks,
     gemini,

pixeltable/functions/anthropic.py

@@ -112,7 +112,7 @@ async def messages(
         to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
 
         >>> msgs = [{'role': 'user', 'content': tbl.prompt}]
-        ... tbl.add_computed_column(response= messages(msgs, model='claude-3-haiku-20240307'))
+        ... tbl.add_computed_column(response=messages(msgs, model='claude-3-haiku-20240307'))
     """
 
     # it doesn't look like count_tokens() actually exists in the current version of the library

pixeltable/functions/bedrock.py

@@ -0,0 +1,130 @@
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+import pixeltable as pxt
+from pixeltable import env, exprs
+from pixeltable.func import Tools
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    from botocore.client import BaseClient
+
+_logger = logging.getLogger('pixeltable')
+
+
+@env.register_client('bedrock')
+def _() -> 'BaseClient':
+    import boto3
+
+    return boto3.client(service_name='bedrock-runtime')
+
+
+# boto3 typing is weird; type information is dynamically defined, so the best we can do for the static checker is `Any`
+def _bedrock_client() -> Any:
+    return env.Env.get().get_client('bedrock')
+
+
+@pxt.udf
+def converse(
+    messages: list[dict[str, Any]],
+    *,
+    model_id: str,
+    system: Optional[list[dict[str, Any]]] = None,
+    inference_config: Optional[dict] = None,
+    additional_model_request_fields: Optional[dict] = None,
+    tool_config: Optional[list[dict]] = None,
+) -> dict:
+    """
+    Generate a conversation response.
+
+    Equivalent to the AWS Bedrock `converse` API endpoint.
+    For additional details, see: <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    __Requirements:__
+
+    - `pip install boto3`
+
+    Args:
+        messages: Input messages.
+        model_id: The model that will complete your prompt.
+        system: An optional system prompt.
+        inference_config: Base inference parameters to use.
+        additional_model_request_fields: Additional inference parameters to use.
+
+    For details on the optional parameters, see:
+    <https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html>
+
+    Returns:
+        A dictionary containing the response and other metadata.
+
+    Examples:
+        Add a computed column that applies the model `anthropic.claude-3-haiku-20240307-v1:0`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> msgs = [{'role': 'user', 'content': [{'text': tbl.prompt}]}]
+        ... tbl.add_computed_column(response=messages(msgs, model_id='anthropic.claude-3-haiku-20240307-v1:0'))
+    """
+
+    kwargs: dict[str, Any] = {'messages': messages, 'modelId': model_id}
+
+    if system is not None:
+        kwargs['system'] = system
+    if inference_config is not None:
+        kwargs['inferenceConfig'] = inference_config
+    if additional_model_request_fields is not None:
+        kwargs['additionalModelRequestFields'] = additional_model_request_fields
+
+    if tool_config is not None:
+        tool_config_ = {
+            'tools': [
+                {
+                    'toolSpec': {
+                        'name': tool['name'],
+                        'description': tool['description'],
+                        'inputSchema': {
+                            'json': {
+                                'type': 'object',
+                                'properties': tool['parameters']['properties'],
+                                'required': tool['required'],
+                            }
+                        },
+                    }
+                }
+                for tool in tool_config
+            ]
+        }
+        kwargs['toolConfig'] = tool_config_
+
+    return _bedrock_client().converse(**kwargs)
+
+
+def invoke_tools(tools: Tools, response: exprs.Expr) -> exprs.InlineDict:
+    """Converts an Anthropic response dict to Pixeltable tool invocation format and calls `tools._invoke()`."""
+    return tools._invoke(_bedrock_response_to_pxt_tool_calls(response))
+
+
+@pxt.udf
+def _bedrock_response_to_pxt_tool_calls(response: dict) -> Optional[dict]:
+    if response.get('stopReason') != 'tool_use':
+        return None
+
+    pxt_tool_calls: dict[str, list[dict[str, Any]]] = {}
+    for message in response['output']['message']['content']:
+        if 'toolUse' in message:
+            tool_call = message['toolUse']
+            tool_name = tool_call['name']
+            if tool_name not in pxt_tool_calls:
+                pxt_tool_calls[tool_name] = []
+            pxt_tool_calls[tool_name].append({'args': tool_call['input']})
+
+    if len(pxt_tool_calls) == 0:
+        return None
+
+    return pxt_tool_calls
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__

pixeltable/functions/date.py

@@ -0,0 +1,185 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `DateType`.
+
+Usage example:
+```python
+import pixeltable as pxt
+
+t = pxt.get_table(...)
+t.select(t.date_col.year, t.date_col.weekday()).collect()
+```
+"""
+
+from datetime import date, timedelta
+
+import sqlalchemy as sql
+
+import pixeltable as pxt
+from pixeltable.utils.code import local_public_names
+
+_SQL_ZERO = sql.literal(0)
+
+# NOT YET SUPPORTED date +/- integer
+# NOT YET SUPPORTED date1 - date2 -> integer
+# NOT YET SUPPORTED timestamp(date)
+# NOT YET SUPPORTED date(timestamp)
+
+
+@pxt.udf(is_property=True)
+def year(self: date) -> int:
+    """
+    Between [`MINYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MINYEAR) and
+    [`MAXYEAR`](https://docs.python.org/3/library/datetime.html#datetime.MAXYEAR) inclusive.
+
+    Equivalent to [`date.year`](https://docs.python.org/3/library/datetime.html#datetime.date.year).
+    """
+    return self.year
+
+
+@year.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('year', self)
+
+
+@pxt.udf(is_property=True)
+def month(self: date) -> int:
+    """
+    Between 1 and 12 inclusive.
+
+    Equivalent to [`date.month`](https://docs.python.org/3/library/datetime.html#datetime.date.month).
+    """
+    return self.month
+
+
+@month.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('month', self)
+
+
+@pxt.udf(is_property=True)
+def day(self: date) -> int:
+    """
+    Between 1 and the number of days in the given month of the given year.
+
+    Equivalent to [`date.day`](https://docs.python.org/3/library/datetime.html#datetime.date.day).
+    """
+    return self.day
+
+
+@day.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('day', self)
+
+
+@pxt.udf(is_method=True)
+def make_date(year: int, month: int, day: int) -> date:
+    """
+    Create a date.
+
+    Equivalent to [`datetime()`](https://docs.python.org/3/library/datetime.html#datetime.date).
+    """
+    return date(year, month, day)
+
+
+@make_date.to_sql
+def _(year: sql.ColumnElement, month: sql.ColumnElement, day: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.func.make_date(sql.cast(year, sql.Integer), sql.cast(month, sql.Integer), sql.cast(day, sql.Integer))
+
+
+@pxt.udf(is_method=True)
+def weekday(self: date) -> int:
+    """
+    Between 0 (Monday) and 6 (Sunday) inclusive.
+
+    Equivalent to [`date.weekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.weekday).
+    """
+    return self.weekday()
+
+
+@weekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self) - 1
+
+
+@pxt.udf(is_method=True)
+def isoweekday(self: date) -> int:
+    """
+    Return the day of the week as an integer, where Monday is 1 and Sunday is 7.
+
+    Equivalent to [`date.isoweekday()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoweekday).
+    """
+    return self.isoweekday()
+
+
+@isoweekday.to_sql
+def _(self: sql.ColumnElement) -> sql.ColumnElement:
+    return sql.extract('isodow', self)
+
+
+@pxt.udf(is_method=True)
+def isocalendar(self: date) -> dict:
+    """
+    Return a dictionary with three entries: `'year'`, `'week'`, and `'weekday'`.
+
+    Equivalent to
+    [`date.isocalendar()`](https://docs.python.org/3/library/datetime.html#datetime.date.isocalendar).
+    """
+    iso_year, iso_week, iso_weekday = self.isocalendar()
+    return {'year': iso_year, 'week': iso_week, 'weekday': iso_weekday}
+
+
+@pxt.udf(is_method=True)
+def isoformat(self: date, sep: str = 'T', timespec: str = 'auto') -> str:
+    """
+    Return a string representing the date and time in ISO 8601 format.
+
+    Equivalent to [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat).
+
+    Args:
+        sep: Separator between date and time.
+        timespec: The number of additional terms in the output. See the
+            [`date.isoformat()`](https://docs.python.org/3/library/datetime.html#datetime.date.isoformat)
+            documentation for more details.
+    """
+    return self.isoformat()
+
+
+@pxt.udf(is_method=True)
+def toordinal(self: date) -> int:
+    """
+    Return the proleptic Gregorian ordinal of the date, where January 1 of year 1 has ordinal 1.
+
+    Equivalent to [`date.toordinal()`](https://docs.python.org/3/library/datetime.html#datetime.date.toordinal).
+    """
+    return self.toordinal()
+
+
+@pxt.udf(is_method=True)
+def strftime(self: date, format: str) -> str:
+    """
+    Return a string representing the date and time, controlled by an explicit format string.
+
+    Equivalent to [`date.strftime()`](https://docs.python.org/3/library/datetime.html#datetime.date.strftime).
+
+    Args:
+        format: The format string to control the output. For a complete list of formatting directives, see
+            [`strftime()` and `strptime()` Behavior](https://docs.python.org/3/library/datetime.html#strftime-and-strptime-behavior).
+    """
+    return self.strftime(format)
+
+
+@pxt.udf(is_method=True)
+def add_days(self: date, n: int) -> date:
+    """
+    Add `n` days to the date.
+
+    Equivalent to [`date + timedelta(days=n)`](https://docs.python.org/3/library/datetime.html#datetime.timedelta).
+    """
+    return self + timedelta(days=n)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__() -> list[str]:
+    return __all__