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

pixeltable/dataframe.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import builtins
 import copy
 import hashlib
 import json
@@ -7,7 +8,7 @@ import logging
 import mimetypes
 import traceback
 from pathlib import Path
-from typing import List, Optional, Any, Dict, Iterator, Tuple, Set, Callable
+from typing import TYPE_CHECKING, Any, Callable, Dict, Hashable, Iterator, List, Optional, Set, Tuple
 
 import pandas as pd
 import pandas.io.formats.style
@@ -16,6 +17,7 @@ import sqlalchemy as sql
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.exprs as exprs
+from pixeltable import exec
 from pixeltable.catalog import is_valid_identifier
 from pixeltable.catalog.globals import UpdateStatus
 from pixeltable.env import Env
@@ -24,6 +26,9 @@ from pixeltable.type_system import ColumnType
 from pixeltable.utils.formatter import Formatter
 from pixeltable.utils.http_server import get_file_uri
 
+if TYPE_CHECKING:
+    import torch
+
 __all__ = ['DataFrame']
 
 _logger = logging.getLogger('pixeltable')
@@ -38,27 +43,25 @@ def _create_source_tag(file_path: str) -> str:
 
 
 class DataFrameResultSet:
-    def __init__(self, rows: List[List[Any]], col_names: List[str], col_types: List[ColumnType]):
+    def __init__(self, rows: list[list[Any]], schema: dict[str, ColumnType]):
         self._rows = rows
-        self._col_names = col_names
-        self._col_types = col_types
+        self._col_names = list(schema.keys())
+        self.__schema = schema
         self.__formatter = Formatter(len(self._rows), len(self._col_names), Env.get().http_address)
 
+    @property
+    def schema(self) -> dict[str, ColumnType]:
+        return self.__schema
+
     def __len__(self) -> int:
         return len(self._rows)
 
-    def column_names(self) -> List[str]:
-        return self._col_names
-
-    def column_types(self) -> List[ColumnType]:
-        return self._col_types
-
     def __repr__(self) -> str:
         return self.to_pandas().__repr__()
 
     def _repr_html_(self) -> str:
-        formatters: dict[str, Callable] = {}
-        for col_name, col_type in zip(self._col_names, self._col_types):
+        formatters: dict[Hashable, Callable[[object], str]] = {}
+        for col_name, col_type in self.schema.items():
             formatter = self.__formatter.get_pandas_formatter(col_type)
             if formatter is not None:
                 formatters[col_name] = formatter
@@ -169,8 +172,9 @@ class DataFrame:
         DataFrame._select_list_check_rep(list(zip(select_list_exprs, column_names)))
         # check select list after expansion to catch early
         # the following two lists are always non empty, even if select list is None.
+        assert len(column_names) == len(select_list_exprs)
         self._select_list_exprs = select_list_exprs
-        self._column_names = column_names
+        self._schema = {column_names[i]: select_list_exprs[i].col_type for i in range(len(column_names))}
         self.select_list = select_list
 
         self.where_clause = copy.deepcopy(where_clause)
@@ -202,22 +206,20 @@ class DataFrame:
     def _normalize_select_list(
         cls,
         tbl: catalog.TableVersionPath,
-        select_list: Optional[List[Tuple[exprs.Expr, Optional[str]]]],
-    ) -> Tuple[List[exprs.Expr], List[str]]:
+        select_list: Optional[list[tuple[exprs.Expr, Optional[str]]]],
+    ) -> tuple[list[exprs.Expr], list[str]]:
         """
         Expand select list information with all columns and their names
         Returns:
             a pair composed of the list of expressions and the list of corresponding names
         """
         if select_list is None:
-            expanded_list = [(exprs.ColumnRef(col), None) for col in tbl.columns()]
-        else:
-            expanded_list = select_list
+            select_list = [(exprs.ColumnRef(col), None) for col in tbl.columns()]
 
-        out_exprs: List[exprs.Expr] = []
-        out_names: List[str] = []  # keep track of order
+        out_exprs: list[exprs.Expr] = []
+        out_names: list[str] = []  # keep track of order
         seen_out_names: set[str] = set()  # use to check for duplicates in loop, avoid square complexity
-        for i, (expr, name) in enumerate(expanded_list):
+        for i, (expr, name) in enumerate(select_list):
             if name is None:
                 # use default, add suffix if needed so default adds no duplicates
                 default_name = expr.default_column_name()
@@ -275,6 +277,24 @@ class DataFrame:
         """Run the query and return rows as a generator.
         This function must not modify the state of the DataFrame, otherwise it breaks dataset caching.
         """
+        plan = self._create_query_plan()
+
+        def exec_plan(conn: sql.engine.Connection) -> Iterator[exec.DataRowBatch]:
+            plan.ctx.set_conn(conn)
+            plan.open()
+            try:
+                for row_batch in plan:
+                    yield from row_batch
+            finally:
+                plan.close()
+
+        if conn is None:
+            with Env.get().engine.begin() as conn:
+                yield from exec_plan(conn)
+        else:
+            yield from exec_plan(conn)
+
+    def _create_query_plan(self) -> exec.ExecNode:
         # construct a group-by clause if we're grouping by a table
         group_by_clause: List[exprs.Expr] = []
         if self.grouping_tbl is not None:
@@ -289,7 +309,7 @@ class DataFrame:
         for item in self._select_list_exprs:
             item.bind_rel_paths(None)
 
-        plan = Planner.create_query_plan(
+        return Planner.create_query_plan(
             self.tbl,
             self._select_list_exprs,
             where_clause=self.where_clause,
@@ -298,21 +318,6 @@ class DataFrame:
             limit=self.limit_val if self.limit_val is not None else 0,
         )  # limit_val == 0: no limit_val
 
-        def exec_plan(conn: sql.engine.Connection) -> Iterator[exprs.DataRow]:
-            plan.ctx.set_conn(conn)
-            plan.open()
-            try:
-                for row_batch in plan:
-                    for data_row in row_batch:
-                        yield data_row
-            finally:
-                plan.close()
-
-        if conn is None:
-            with Env.get().engine.begin() as conn:
-                yield from exec_plan(conn)
-        else:
-            yield from exec_plan(conn)
 
     def show(self, n: int = 20) -> DataFrameResultSet:
         assert n is not None
@@ -334,11 +339,9 @@ class DataFrame:
         result._reverse()
         return result
 
-    def get_column_names(self) -> List[str]:
-        return self._column_names
-
-    def get_column_types(self) -> List[ColumnType]:
-        return [expr.col_type for expr in self._select_list_exprs]
+    @property
+    def schema(self) -> dict[str, ColumnType]:
+        return self._schema
 
     def bind(self, args: dict[str, Any]) -> DataFrame:
         """Bind arguments to parameters and return a new DataFrame."""
@@ -369,7 +372,7 @@ class DataFrame:
         if order_by_exprs is not None:
             exprs.Expr.list_substitute(order_by_exprs, var_exprs)
 
-        select_list = list(zip(select_list_exprs, self._column_names))
+        select_list = list(zip(select_list_exprs, self.schema.keys()))
         order_by_clause: Optional[list[tuple[exprs.Expr, bool]]] = None
         if order_by_exprs is not None:
             order_by_clause = [
@@ -409,8 +412,7 @@ class DataFrame:
         except sql.exc.DBAPIError as e:
             raise excs.Error(f'Error during SQL execution:\n{e}')
 
-        col_types = self.get_column_types()
-        return DataFrameResultSet(result_rows, self._column_names, col_types)
+        return DataFrameResultSet(result_rows, self.schema)
 
     def count(self) -> int:
         from pixeltable.plan import Planner
@@ -429,7 +431,7 @@ class DataFrame:
             assert len(self.select_list) > 0
             heading_vals.append('Select')
             heading_vals.extend([''] * (len(self.select_list) - 1))
-            info_vals.extend(self.get_column_names())
+            info_vals.extend(self.schema.keys())
         if self.where_clause is not None:
             heading_vals.append('Where')
             info_vals.append(self.where_clause.display_str(inline=False))
@@ -457,7 +459,7 @@ class DataFrame:
         # white-space: pre-wrap: print \n as newline
         # th: center-align headings
         return (
-            pd_df.style.set_properties(**{'white-space': 'pre-wrap', 'text-align': 'left'})
+            pd_df.style.set_properties(None, **{'white-space': 'pre-wrap', 'text-align': 'left'})
             .set_table_styles([dict(selector='th', props=[('text-align', 'center')])])
             .hide(axis='index')
             .hide(axis='columns')
@@ -469,19 +471,17 @@ class DataFrame:
         The description has two columns, heading and info, which list the contents of each 'component'
                 (select list, where clause, ...) vertically.
         """
-        try:
-            __IPYTHON__
+        if getattr(builtins, '__IPYTHON__', False):
             from IPython.display import display
-
             display(self._description_html())
-        except NameError:
+        else:
             print(self.__repr__())
 
     def __repr__(self) -> str:
         return self._description().to_string(header=False, index=False)
 
     def _repr_html_(self) -> str:
-        return self._description_html()._repr_html_()
+        return self._description_html()._repr_html_()  # type: ignore[attr-defined]
 
     def select(self, *items: Any, **named_items: Any) -> DataFrame:
         if self.select_list is not None:
@@ -562,7 +562,7 @@ class DataFrame:
                 # we need to make sure that the grouping table is a base of self.tbl
                 base = self.tbl.find_tbl_version(item._tbl_version_path.tbl_id())
                 if base is None or base.id == self.tbl.tbl_id():
-                    raise excs.Error(f'group_by(): {item.name} is not a base table of {self.tbl.tbl_name()}')
+                    raise excs.Error(f'group_by(): {item._name} is not a base table of {self.tbl.tbl_name()}')
                 grouping_tbl = item._tbl_version_path.tbl_version
                 break
             if not isinstance(item, exprs.Expr):
@@ -756,12 +756,12 @@ class DataFrame:
         Env.get().require_package('torch')
         Env.get().require_package('torchvision')
 
-        from pixeltable.io.parquet import save_parquet  # pylint: disable=import-outside-toplevel
-        from pixeltable.utils.pytorch import PixeltablePytorchDataset  # pylint: disable=import-outside-toplevel
+        from pixeltable.io.parquet import save_parquet
+        from pixeltable.utils.pytorch import PixeltablePytorchDataset
 
         cache_key = self._hash_result_set()
 
-        dest_path = (Env.get().dataset_cache_dir / f'df_{cache_key}').with_suffix('.parquet')  # pylint: disable = protected-access
+        dest_path = (Env.get().dataset_cache_dir / f'df_{cache_key}').with_suffix('.parquet')
         if dest_path.exists():  # fast path: use cache
             assert dest_path.is_dir()
         else:

pixeltable/env.py

@@ -268,7 +268,7 @@ class Env:
 
         # in pixeltable_pgserver.get_server(): cleanup_mode=None will leave db on for debugging purposes
         self._db_server = pixeltable_pgserver.get_server(self._pgdata_dir, cleanup_mode=None)
-        self._db_url = self._db_server.get_uri(database=self._db_name)
+        self._db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
 
         if reinit_db:
             if self._store_db_exists():
@@ -297,7 +297,7 @@ class Env:
     def _store_db_exists(self) -> bool:
         assert self._db_name is not None
         # don't try to connect to self.db_name, it may not exist
-        db_url = self._db_server.get_uri(database='postgres')
+        db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
         engine = sql.create_engine(db_url, future=True)
         try:
             with engine.begin() as conn:
@@ -312,7 +312,7 @@ class Env:
     def _create_store_db(self) -> None:
         assert self._db_name is not None
         # create the db
-        pg_db_url = self._db_server.get_uri(database='postgres')
+        pg_db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
         engine = sql.create_engine(pg_db_url, future=True, isolation_level='AUTOCOMMIT')
         preparer = engine.dialect.identifier_preparer
         try:
@@ -327,7 +327,7 @@ class Env:
             engine.dispose()
 
         # enable pgvector
-        store_db_url = self._db_server.get_uri(database=self._db_name)
+        store_db_url = self._db_server.get_uri(database=self._db_name, driver='psycopg')
         engine = sql.create_engine(store_db_url, future=True, isolation_level='AUTOCOMMIT')
         try:
             with engine.begin() as conn:
@@ -337,7 +337,7 @@ class Env:
 
     def _drop_store_db(self) -> None:
         assert self._db_name is not None
-        db_url = self._db_server.get_uri(database='postgres')
+        db_url = self._db_server.get_uri(database='postgres', driver='psycopg')
         engine = sql.create_engine(db_url, future=True, isolation_level='AUTOCOMMIT')
         preparer = engine.dialect.identifier_preparer
         try:
@@ -425,6 +425,7 @@ class Env:
             else:
                 self._installed_packages[package] = None
 
+        check('toml')
         check('datasets')
         check('torch')
         check('torchvision')
@@ -443,6 +444,7 @@ class Env:
             self._spacy_nlp = spacy.load('en_core_web_sm')
         check('tiktoken')
         check('openai')
+        check('anthropic')
         check('together')
         check('fireworks')
         check('label_studio_sdk')

pixeltable/exec/__init__.py

@@ -1,11 +1,11 @@
 from .aggregation_node import AggregationNode
 from .cache_prefetch_node import CachePrefetchNode
 from .component_iteration_node import ComponentIterationNode
+from .data_row_batch import DataRowBatch
 from .exec_context import ExecContext
 from .exec_node import ExecNode
 from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
-from .sql_node import SqlScanNode, SqlLookupNode
-from .row_update_node import RowUpdateNode
 from .media_validation_node import MediaValidationNode
-from .data_row_batch import DataRowBatch
+from .row_update_node import RowUpdateNode
+from .sql_node import SqlLookupNode, SqlScanNode

pixeltable/exec/aggregation_node.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 import sys
-from typing import List, Optional, Any
+from typing import Iterable, List, Optional, Any
 
 import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
@@ -15,12 +15,12 @@ _logger = logging.getLogger('pixeltable')
 class AggregationNode(ExecNode):
     def __init__(
             self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, group_by: List[exprs.Expr],
-            agg_fn_calls: List[exprs.FunctionCall], input_exprs: List[exprs.Expr], input: ExecNode
+            agg_fn_calls: List[exprs.FunctionCall], input_exprs: Iterable[exprs.Expr], input: ExecNode
     ):
         super().__init__(row_builder, group_by + agg_fn_calls, input_exprs, input)
         self.input = input
         self.group_by = group_by
-        self.input_exprs = input_exprs
+        self.input_exprs = list(input_exprs)
         self.agg_fn_calls = agg_fn_calls
         self.agg_fn_eval_ctx = row_builder.create_eval_ctx(agg_fn_calls, exclude=input_exprs)
         self.output_batch = DataRowBatch(tbl, row_builder, 0)

pixeltable/exec/expr_eval_node.py

@@ -3,7 +3,7 @@ import sys
 import time
 import warnings
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Iterable, List, Optional
 
 from tqdm import tqdm, TqdmWarning
 
@@ -23,12 +23,12 @@ class ExprEvalNode(ExecNode):
         """List of exprs that form an evaluation context and contain calls to at most one external function"""
         exprs: List[exprs.Expr]
         batched_fn: Optional[CallableFunction]
-        segment_ctxs: List[exprs.RowBuilder.EvalCtx]
+        segment_ctxs: List['exprs.RowBuilder.EvalCtx']
         target_slot_idxs: List[int]
         batch_size: int = 8
 
     def __init__(
-            self, row_builder: exprs.RowBuilder, output_exprs: List[exprs.Expr], input_exprs: List[exprs.Expr],
+            self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr], input_exprs: Iterable[exprs.Expr],
             input: ExecNode
     ):
         super().__init__(row_builder, output_exprs, input_exprs, input)

pixeltable/exec/in_memory_data_node.py

@@ -1,9 +1,10 @@
 import logging
-from typing import List, Dict, Any, Optional
+from typing import Any, Optional
 
 import pixeltable.catalog as catalog
 import pixeltable.exprs as exprs
 from pixeltable.utils.media_store import MediaStore
+
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
 
@@ -18,8 +19,8 @@ class InMemoryDataNode(ExecNode):
     - if an input row doesn't provide a value, sets the slot to the column default
     """
     def __init__(
-            self, tbl: catalog.TableVersionPath, rows: List[Dict[str, Any]],
-            row_builder: exprs.RowBuilder, start_row_id: int,
+        self, tbl: catalog.TableVersion, rows: list[dict[str, Any]],
+        row_builder: exprs.RowBuilder, start_row_id: int,
     ):
         # we materialize all output slots
         output_exprs = [e for e in row_builder.get_output_exprs() if isinstance(e, exprs.ColumnRef)]
@@ -75,4 +76,3 @@ class InMemoryDataNode(ExecNode):
         self.has_returned_data = True
         _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_rows)} output_rows')
         return self.output_rows
-

pixeltable/exec/sql_node.py

@@ -258,6 +258,10 @@ class SqlLookupNode(SqlNode):
     """
     Materializes data from the store via a Select stmt with a WHERE clause that matches a list of key values
     """
+
+    stmt: sql.Select
+    where_clause: sql.ColumnElement[bool]
+
     def __init__(
             self, tbl: catalog.TableVersionPath, row_builder: exprs.RowBuilder,
             select_list: Iterable[exprs.Expr], sa_key_cols: list[sql.Column], key_vals: list[tuple],
@@ -287,4 +291,3 @@ class SqlLookupNode(SqlNode):
             _logger.debug(f'SqlLookupNode stmt:\n{stmt_str}')
         except Exception as e:
             pass
-

pixeltable/exprs/array_slice.py

@@ -1,14 +1,13 @@
 from __future__ import annotations
-from typing import Optional, List, Any, Dict, Tuple
-import copy
+
+from typing import Any, Dict, List, Optional, Tuple
 
 import sqlalchemy as sql
 
+from .data_row import DataRow
 from .expr import Expr
 from .globals import print_slice
-from .data_row import DataRow
 from .row_builder import RowBuilder
-import pixeltable.catalog as catalog
 
 
 class ArraySlice(Expr):

pixeltable/exprs/column_ref.py

@@ -19,19 +19,29 @@ class ColumnRef(Expr):
     For that reason, a ColumnRef needs to be serialized with the qualifying table id (column ids are only
     unique in the context of a particular table).
     """
+
+    col: catalog.Column
+    is_unstored_iter_col: bool
+    iter_arg_ctx: Optional[RowBuilder.EvalCtx]
+    base_rowid_len: int
+    base_rowid: list[Optional[Any]]
+    iterator: Optional[iters.ComponentIterator]
+    pos_idx: Optional[int]
+    id: int
+
     def __init__(self, col: catalog.Column):
         super().__init__(col.col_type)
         assert col.tbl is not None
         self.col = col
         self.is_unstored_iter_col = \
             col.tbl.is_component_view() and col.tbl.is_iterator_column(col) and not col.is_stored
-        self.iter_arg_ctx: Optional[RowBuilder.EvalCtx] = None
+        self.iter_arg_ctx = None
         # number of rowid columns in the base table
         self.base_rowid_len = col.tbl.base.num_rowid_columns() if self.is_unstored_iter_col else 0
         self.base_rowid = [None] * self.base_rowid_len
-        self.iterator: Optional[iters.ComponentIterator] = None
+        self.iterator = None
         # index of the position column in the view's primary key; don't try to reference tbl.store_tbl here
-        self.pos_idx: Optional[int] = col.tbl.num_rowid_columns() - 1 if self.is_unstored_iter_col else None
+        self.pos_idx = col.tbl.num_rowid_columns() - 1 if self.is_unstored_iter_col else None
         self.id = self._create_id()
 
     def set_iter_arg_ctx(self, iter_arg_ctx: RowBuilder.EvalCtx) -> None:
@@ -74,7 +84,13 @@ class ColumnRef(Expr):
         return self.col == other.col
 
     def __str__(self) -> str:
-        return self.col.name
+        if self.col.name is None:
+            return f'<unnamed column {self.col.id}>'
+        else:
+            return self.col.name
+
+    def __repr__(self) -> str:
+        return f'ColumnRef({self.col!r})'
 
     def sql_expr(self) -> Optional[sql.ClauseElement]:
         return self.col.sa_col

pixeltable/exprs/comparison.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from typing import Optional, List, Any, Dict
 
 import sqlalchemy as sql
@@ -78,6 +79,7 @@ class Comparison(Expr):
         right = self._op2.sql_expr()
         if left is None or right is None:
             return None
+
         if self.operator == ComparisonOperator.LT:
             return left < right
         if self.operator == ComparisonOperator.LE:
@@ -92,18 +94,21 @@ class Comparison(Expr):
             return left >= right
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        left = data_row[self._op1.slot_idx]
+        right = data_row[self._op2.slot_idx]
+
         if self.operator == ComparisonOperator.LT:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] < data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left < right
         elif self.operator == ComparisonOperator.LE:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] <= data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left <= right
         elif self.operator == ComparisonOperator.EQ:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] == data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left == right
         elif self.operator == ComparisonOperator.NE:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] != data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left != right
         elif self.operator == ComparisonOperator.GT:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] > data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left > right
         elif self.operator == ComparisonOperator.GE:
-            data_row[self.slot_idx] = data_row[self._op1.slot_idx] >= data_row[self._op2.slot_idx]
+            data_row[self.slot_idx] = left >= right
 
     def _as_dict(self) -> Dict:
         return {'operator': self.operator.value, **super()._as_dict()}

pixeltable/exprs/data_row.py

@@ -96,6 +96,9 @@ class DataRow:
             self.file_paths[slot_idx] = None
             self.file_urls[slot_idx] = None
 
+    def __len__(self) -> int:
+        return len(self.vals)
+
     def __getitem__(self, index: object) -> Any:
         """Returns in-memory value, ie, what is needed for expr evaluation"""
         if not self.has_val[index]: