pixeltable 0.2.5__py3-none-any.whl → 0.2.7__py3-none-any.whl

pixeltable/env.py

@@ -5,39 +5,48 @@ import glob
 import http.server
 import importlib
 import importlib.util
+import inspect
 import logging
 import os
-import socketserver
 import sys
 import threading
 import uuid
 import warnings
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, Optional, Dict, Any, List
 
 import pgserver
 import sqlalchemy as sql
 import yaml
-from sqlalchemy_utils.functions import database_exists, create_database, drop_database
 from tqdm import TqdmWarning
 
 import pixeltable.exceptions as excs
 from pixeltable import metadata
+from pixeltable.utils.http_server import make_server
 
 
 class Env:
     """
     Store for runtime globals.
     """
+
     _instance: Optional[Env] = None
     _log_fmt_str = '%(asctime)s %(levelname)s %(name)s %(filename)s:%(lineno)d: %(message)s'
 
     @classmethod
     def get(cls) -> Env:
         if cls._instance is None:
-            cls._instance = Env()
+            cls._init_env()
         return cls._instance
 
+    @classmethod
+    def _init_env(cls, reinit_db: bool = False) -> None:
+        env = Env()
+        env._set_up(reinit_db=reinit_db)
+        env._upgrade_metadata()
+        cls._instance = env
+
     def __init__(self):
         self._home: Optional[Path] = None
         self._media_dir: Optional[Path] = None  # computed media files
@@ -46,7 +55,7 @@ class Env:
         self._log_dir: Optional[Path] = None  # log files
         self._tmp_dir: Optional[Path] = None  # any tmp files
         self._sa_engine: Optional[sql.engine.base.Engine] = None
-        self._pgdata_dir : Optional[Path] = None
+        self._pgdata_dir: Optional[Path] = None
         self._db_name: Optional[str] = None
         self._db_server: Optional[pgserver.PostgresServer] = None
         self._db_url: Optional[str] = None
@@ -54,12 +63,11 @@ class Env:
         # info about installed packages that are utilized by some parts of the code;
         # package name -> version; version == []: package is installed, but we haven't determined the version yet
         self._installed_packages: Dict[str, Optional[List[int]]] = {}
-        self._nos_client: Optional[Any] = None
         self._spacy_nlp: Optional[Any] = None  # spacy.Language
-        self._httpd: Optional[socketserver.TCPServer] = None
+        self._httpd: Optional[http.server.HTTPServer] = None
         self._http_address: Optional[str] = None
 
-        self._registered_clients: dict[str, Any] = {}
+        self._registered_clients: dict[str, ApiClient] = {}
 
         # logging-related state
         self._logger = logging.getLogger('pixeltable')
@@ -94,13 +102,43 @@ class Env:
         assert self._http_address is not None
         return self._http_address
 
+    def configure_logging(
+        self,
+        *,
+        to_stdout: Optional[bool] = None,
+        level: Optional[int] = None,
+        add: Optional[str] = None,
+        remove: Optional[str] = None,
+    ) -> None:
+        """Configure logging.
+
+        Args:
+            to_stdout: if True, also log to stdout
+            level: default log level
+            add: comma-separated list of 'module name:log level' pairs; ex.: add='video:10'
+            remove: comma-separated list of module names
+        """
+        if to_stdout is not None:
+            self.log_to_stdout(to_stdout)
+        if level is not None:
+            self.set_log_level(level)
+        if add is not None:
+            for module, level_str in [t.split(':') for t in add.split(',')]:
+                self.set_module_log_level(module, int(level_str))
+        if remove is not None:
+            for module in remove.split(','):
+                self.set_module_log_level(module, None)
+        if to_stdout is None and level is None and add is None and remove is None:
+            self.print_log_config()
+
     def print_log_config(self) -> None:
         print(f'logging to {self._logfilename}')
         print(f'{"" if self._log_to_stdout else "not "}logging to stdout')
         print(f'default log level: {logging.getLevelName(self._default_log_level)}')
         print(
             f'module log levels: '
-            f'{",".join([name + ":" + logging.getLevelName(val) for name, val in self._module_log_level.items()])}')
+            f'{",".join([name + ":" + logging.getLevelName(val) for name, val in self._module_log_level.items()])}'
+        )
 
     def log_to_stdout(self, enable: bool = True) -> None:
         self._log_to_stdout = enable
@@ -135,10 +173,14 @@ class Env:
         else:
             return False
 
-    def set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
+    def _set_up(self, echo: bool = False, reinit_db: bool = False) -> None:
         if self._initialized:
             return
 
+        # Disable spurious warnings
+        warnings.simplefilter('ignore', category=TqdmWarning)
+        os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
         self._initialized = True
         home = Path(os.environ.get('PIXELTABLE_HOME', str(Path.home() / '.pixeltable')))
         assert self._home is None or self._home == home
@@ -204,6 +246,14 @@ class Env:
         av_logger.addHandler(av_fh)
         av_logger.propagate = False
 
+        # configure web-server logging
+        http_logfilename = self._logfilename.replace('.log', '_http.log')
+        http_fh = logging.FileHandler(self._log_dir / http_logfilename, mode='w')
+        http_fh.setFormatter(logging.Formatter(self._log_fmt_str))
+        http_logger = logging.getLogger('pixeltable.http.server')
+        http_logger.addHandler(http_fh)
+        http_logger.propagate = False
+
         # empty tmp dir
         for path in glob.glob(f'{self._tmp_dir}/*'):
             os.remove(path)
@@ -216,23 +266,19 @@ class Env:
         self._db_url = self._db_server.get_uri(database=self._db_name)
 
         if reinit_db:
-            if database_exists(self.db_url):
-                drop_database(self.db_url)
+            if self._store_db_exists():
+                self._drop_store_db()
 
-        if not database_exists(self.db_url):
+        if not self._store_db_exists():
             self._logger.info(f'creating database at {self.db_url}')
-            create_database(self.db_url)
-            self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+            self._create_store_db()
+            self._create_engine(echo=echo)
             from pixeltable.metadata import schema
             schema.Base.metadata.create_all(self._sa_engine)
             metadata.create_system_info(self._sa_engine)
-            # enable pgvector
-            with self._sa_engine.begin() as conn:
-                conn.execute(sql.text('CREATE EXTENSION vector'))
         else:
             self._logger.info(f'found database {self.db_url}')
-            if self._sa_engine is None:
-                self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True)
+            self._create_engine(echo=echo)
 
         print(f'Connected to Pixeltable database at: {self.db_url}')
 
@@ -240,72 +286,122 @@ class Env:
         self._set_up_runtime()
         self.log_to_stdout(False)
 
-        # Disable spurious warnings
-        warnings.simplefilter("ignore", category=TqdmWarning)
+    def _create_engine(self, echo: bool = False) -> None:
+        self._sa_engine = sql.create_engine(self.db_url, echo=echo, future=True, isolation_level='AUTOCOMMIT')
 
-    def upgrade_metadata(self) -> None:
+    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')
+        engine = sql.create_engine(db_url, future=True)
+        try:
+            with engine.begin() as conn:
+                stmt = f"SELECT COUNT(*) FROM pg_database WHERE datname = '{self._db_name}'"
+                result = conn.scalar(sql.text(stmt))
+                assert result <= 1
+                return result == 1
+        finally:
+            engine.dispose()
+
+
+    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')
+        engine = sql.create_engine(pg_db_url, future=True, isolation_level='AUTOCOMMIT')
+        preparer = engine.dialect.identifier_preparer
+        try:
+            with engine.begin() as conn:
+                # use C collation to get standard C/Python-style sorting
+                stmt = (
+                    f"CREATE DATABASE {preparer.quote(self._db_name)} "
+                    "ENCODING 'utf-8' LC_COLLATE 'C' LC_CTYPE 'C' TEMPLATE template0"
+                )
+                conn.execute(sql.text(stmt))
+        finally:
+            engine.dispose()
+
+        # enable pgvector
+        store_db_url = self._db_server.get_uri(database=self._db_name)
+        engine = sql.create_engine(store_db_url, future=True, isolation_level='AUTOCOMMIT')
+        try:
+            with engine.begin() as conn:
+                conn.execute(sql.text('CREATE EXTENSION vector'))
+        finally:
+            engine.dispose()
+
+    def _drop_store_db(self) -> None:
+        assert self._db_name is not None
+        db_url = self._db_server.get_uri(database='postgres')
+        engine = sql.create_engine(db_url, future=True, isolation_level='AUTOCOMMIT')
+        preparer = engine.dialect.identifier_preparer
+        try:
+            with engine.begin() as conn:
+                # terminate active connections
+                stmt = (f"""
+                    SELECT pg_terminate_backend(pg_stat_activity.pid)
+                    FROM pg_stat_activity
+                    WHERE pg_stat_activity.datname = '{self._db_name}'
+                    AND pid <> pg_backend_pid()
+                """)
+                conn.execute(sql.text(stmt))
+                # drop db
+                stmt = f'DROP DATABASE {preparer.quote(self._db_name)}'
+                conn.execute(sql.text(stmt))
+        finally:
+            engine.dispose()
+
+    def _upgrade_metadata(self) -> None:
         metadata.upgrade_md(self._sa_engine)
 
-    def _create_nos_client(self) -> None:
-        import nos
-        self._logger.info('connecting to NOS')
-        nos.init(logging_level=logging.DEBUG)
-        self._nos_client = nos.client.InferenceClient()
-        self._logger.info('waiting for NOS')
-        self._nos_client.WaitForServer()
-
-        # now that we have a client, we can create the module
-        import importlib
-        try:
-            importlib.import_module('pixeltable.functions.nos')
-            # it's already been created
-            return
-        except ImportError:
-            pass
-        from pixeltable.functions.util import create_nos_modules
-        _ = create_nos_modules()
+    def _register_client(self, name: str, init_fn: Callable) -> None:
+        sig = inspect.signature(init_fn)
+        param_names = list(sig.parameters.keys())
+        self._registered_clients[name] = ApiClient(init_fn=init_fn, param_names=param_names)
 
-    def get_client(self, name: str, init: Callable, environ: Optional[str] = None) -> Any:
+    def get_client(self, name: str) -> Any:
         """
-        Gets the client with the specified name, using `init` to construct one if necessary.
+        Gets the client with the specified name, initializing it if necessary.
 
-        - name: The name of the client
-        - init: A `Callable` with signature `fn(api_key: str) -> Any` that constructs a client object
-        - environ: The name of the environment variable to use for the API key, if no API key is found in config
-            (defaults to f'{name.upper()}_API_KEY')
+        Args:
+            - name: The name of the client
         """
-        if name in self._registered_clients:
-            return self._registered_clients[name]
-
-        if environ is None:
-            environ = f'{name.upper()}_API_KEY'
-
-        if name in self._config and 'api_key' in self._config[name]:
-            api_key = self._config[name]['api_key']
-        else:
-            api_key = os.environ.get(environ)
-        if api_key is None or api_key == '':
-            raise excs.Error(f'`{name}` client not initialized (no API key configured).')
-
-        client = init(api_key)
-        self._registered_clients[name] = client
+        cl = self._registered_clients[name]
+        if cl.client_obj is not None:
+            return cl.client_obj  # Already initialized
+
+        # Construct a client. For each client parameter, first check if the parameter is in the environment;
+        # if not, look in Pixeltable config from `config.yaml`.
+
+        init_kwargs: dict[str, str] = {}
+        for param in cl.param_names:
+            environ = f'{name.upper()}_{param.upper()}'
+            if environ in os.environ:
+                init_kwargs[param] = os.environ[environ]
+            elif name.lower() in self._config and param in self._config[name.lower()]:
+                init_kwargs[param] = self._config[name.lower()][param.lower()]
+            if param not in init_kwargs or init_kwargs[param] == '':
+                raise excs.Error(
+                    f'`{name}` client not initialized: parameter `{param}` is not configured.\n'
+                    f'To fix this, specify the `{environ}` environment variable, or put `{param.lower()}` in '
+                    f'the `{name.lower()}` section of $PIXELTABLE_HOME/config.yaml.'
+                )
+
+        cl.client_obj = cl.init_fn(**init_kwargs)
         self._logger.info(f'Initialized `{name}` client.')
-        return client
+        return cl.client_obj
 
     def _start_web_server(self) -> None:
         """
         The http server root is the file system root.
         eg: /home/media/foo.mp4 is located at http://127.0.0.1:{port}/home/media/foo.mp4
+        in windows, the server will translate paths like http://127.0.0.1:{port}/c:/media/foo.mp4
         This arrangement enables serving media hosted within _home,
         as well as external media inserted into pixeltable or produced by pixeltable.
         The port is chosen dynamically to prevent conflicts.
         """
         # Port 0 means OS picks one for us.
-        address = ("127.0.0.1", 0)
-        class FixedRootHandler(http.server.SimpleHTTPRequestHandler):
-            def __init__(self, *args, **kwargs):
-                super().__init__(*args, directory='/', **kwargs)
-        self._httpd = socketserver.TCPServer(address, FixedRootHandler)
+        self._httpd = make_server('127.0.0.1', 0)
         port = self._httpd.server_address[1]
         self._http_address = f'http://127.0.0.1:{port}'
 
@@ -335,19 +431,21 @@ class Env:
         check('transformers')
         check('sentence_transformers')
         check('yolox')
+        check('whisperx')
         check('boto3')
+        check('fitz')  # pymupdf
         check('pyarrow')
         check('spacy')  # TODO: deal with en-core-web-sm
         if self.is_installed_package('spacy'):
             import spacy
+
             self._spacy_nlp = spacy.load('en_core_web_sm')
         check('tiktoken')
         check('openai')
         check('together')
         check('fireworks')
-        check('nos')
-        if self.is_installed_package('nos'):
-            self._create_nos_client()
+        check('label_studio_sdk')
+        check('openpyxl')
 
     def require_package(self, package: str, min_version: Optional[List[int]] = None) -> None:
         assert package in self._installed_packages
@@ -357,7 +455,7 @@ class Env:
             return
 
         # check whether we have a version >= the required one
-        if self._installed_packages[package] == []:
+        if not self._installed_packages[package]:
             m = importlib.import_module(package)
             module_version = [int(x) for x in m.__version__.split('.')]
             self._installed_packages[package] = module_version
@@ -365,9 +463,12 @@ class Env:
         if len(min_version) < len(installed_version):
             normalized_min_version = min_version + [0] * (len(installed_version) - len(min_version))
         if any([a < b for a, b in zip(installed_version, normalized_min_version)]):
-            raise excs.Error((
-                f'The installed version of package {package} is {".".join([str[v] for v in installed_version])}, '
-                f'but version  >={".".join([str[v] for v in min_version])} is required'))
+            raise excs.Error(
+                (
+                    f'The installed version of package {package} is {".".join(str(v) for v in installed_version)}, '
+                    f'but version  >={".".join(str(v) for v in min_version)} is required'
+                )
+            )
 
     def num_tmp_files(self) -> int:
         return len(glob.glob(f'{self._tmp_dir}/*'))
@@ -405,11 +506,44 @@ class Env:
         assert self._sa_engine is not None
         return self._sa_engine
 
-    @property
-    def nos_client(self) -> Any:
-        return self._nos_client
-
     @property
     def spacy_nlp(self) -> Any:
         assert self._spacy_nlp is not None
-        return self._spacy_nlp
+        return self._spacy_nlp
+
+
+def register_client(name: str) -> Callable:
+    """Decorator that registers a third-party API client for use by Pixeltable.
+
+    The decorated function is an initialization wrapper for the client, and can have
+    any number of string parameters, with a signature such as:
+
+    ```
+    def my_client(api_key: str, url: str) -> my_client_sdk.Client:
+        return my_client_sdk.Client(api_key=api_key, url=url)
+    ```
+
+    The initialization wrapper will not be called immediately; initialization will
+    be deferred until the first time the client is used. At initialization time,
+    Pixeltable will attempt to load the client parameters from config. For each
+    config parameter:
+    - If an environment variable named MY_CLIENT_API_KEY (for example) is set, use it;
+    - Otherwise, look for 'api_key' in the 'my_client' section of config.yaml.
+
+    If all config parameters are found, Pixeltable calls the initialization function;
+    otherwise it throws an exception.
+
+    Args:
+        - name (str): The name of the API client (e.g., 'openai' or 'label-studio').
+    """
+    def decorator(fn: Callable) -> None:
+        Env.get()._register_client(name, fn)
+
+    return decorator
+
+
+@dataclass
+class ApiClient:
+    init_fn: Callable
+    param_names: list[str]
+    client_obj: Optional[Any] = None

pixeltable/exec/__init__.py

@@ -6,4 +6,5 @@ from .exec_node import ExecNode
 from .expr_eval_node import ExprEvalNode
 from .in_memory_data_node import InMemoryDataNode
 from .sql_scan_node import SqlScanNode
-from .media_validation_node import MediaValidationNode
+from .media_validation_node import MediaValidationNode
+from .data_row_batch import DataRowBatch

pixeltable/exec/data_row_batch.py

@@ -14,9 +14,8 @@ class DataRowBatch:
 
     Contains the metadata needed to initialize DataRows.
     """
-    def __init__(self, tbl: catalog.TableVersion, row_builder: exprs.RowBuilder, len: int = 0):
-        self.tbl_id = tbl.id
-        self.tbl_version = tbl.version
+    def __init__(self, tbl: Optional[catalog.TableVersion], row_builder: exprs.RowBuilder, len: int = 0):
+        self.tbl = tbl
         self.row_builder = row_builder
         self.img_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_image_type()]
         # non-image media slots
@@ -42,9 +41,10 @@ class DataRowBatch:
 
     def set_row_ids(self, row_ids: List[int]) -> None:
         """Sets pks for rows in batch"""
+        assert self.tbl is not None
         assert len(row_ids) == len(self.rows)
         for row, row_id in zip(self.rows, row_ids):
-            row.set_pk((row_id, self.tbl_version))
+            row.set_pk((row_id, self.tbl))
 
     def __len__(self) -> int:
         return len(self.rows)
@@ -57,6 +57,7 @@ class DataRowBatch:
             flushed_slot_idxs: Optional[List[int]] = None
     ) -> None:
         """Flushes images in the given range of rows."""
+        assert self.tbl is not None
         if stored_img_info is None:
             stored_img_info = []
         if flushed_slot_idxs is None:
@@ -67,12 +68,10 @@ class DataRowBatch:
             idx_range = slice(0, len(self.rows))
         for row in self.rows[idx_range]:
             for info in stored_img_info:
-                filepath = str(MediaStore.prepare_media_path(self.tbl_id, info.col.id, self.tbl_version))
+                filepath = str(MediaStore.prepare_media_path(self.tbl.id, info.col.id, self.tbl.version))
                 row.flush_img(info.slot_idx, filepath)
             for slot_idx in flushed_slot_idxs:
                 row.flush_img(slot_idx)
-        #_logger.debug(
-            #f'flushed images in range {idx_range}: slot_idxs={flushed_slot_idxs} stored_img_info={stored_img_info}')
 
     def __iter__(self) -> Iterator[exprs.DataRow]:
         return DataRowBatchIterator(self)

pixeltable/exec/expr_eval_node.py

@@ -1,20 +1,20 @@
-import sys
-import warnings
-from typing import List, Optional, Tuple
-from dataclasses import dataclass, field
 import logging
+import sys
 import time
+import warnings
+from dataclasses import dataclass
+from typing import List, Optional
 
 from tqdm import tqdm, TqdmWarning
 
+import pixeltable.exprs as exprs
+from pixeltable.func import CallableFunction
 from .data_row_batch import DataRowBatch
 from .exec_node import ExecNode
-import pixeltable.exprs as exprs
-import pixeltable.func as func
-
 
 _logger = logging.getLogger('pixeltable')
 
+
 class ExprEvalNode(ExecNode):
     """Materializes expressions
     """
@@ -22,7 +22,7 @@ class ExprEvalNode(ExecNode):
     class Cohort:
         """List of exprs that form an evaluation context and contain calls to at most one external function"""
         exprs: List[exprs.Expr]
-        ext_function: Optional[func.BatchedFunction]
+        batched_fn: Optional[CallableFunction]
         segment_ctxs: List[exprs.RowBuilder.EvalCtx]
         target_slot_idxs: List[int]
         batch_size: int = 8
@@ -63,12 +63,12 @@ class ExprEvalNode(ExecNode):
         if self.pbar is not None:
             self.pbar.close()
 
-    def _get_batched_fn(self, expr: exprs.Expr) -> Optional[func.BatchedFunction]:
-        if not isinstance(expr, exprs.FunctionCall):
-            return None
-        return expr.fn if isinstance(expr.fn, func.BatchedFunction) else None
+    def _get_batched_fn(self, expr: exprs.Expr) -> Optional[CallableFunction]:
+        if isinstance(expr, exprs.FunctionCall) and isinstance(expr.fn, CallableFunction) and expr.fn.is_batched:
+            return expr.fn
+        return None
 
-    def _is_ext_call(self, expr: exprs.Expr) -> bool:
+    def _is_batched_fn_call(self, expr: exprs.Expr) -> bool:
         return self._get_batched_fn(expr) is not None
 
     def _create_cohorts(self) -> None:
@@ -76,14 +76,14 @@ class ExprEvalNode(ExecNode):
         # break up all_exprs into cohorts such that each cohort contains calls to at most one external function;
         # seed the cohorts with only the ext fn calls
         cohorts: List[List[exprs.Expr]] = []
-        current_ext_function: Optional[func.BatchedFunction] = None
+        current_batched_fn: Optional[CallableFunction] = None
         for e in all_exprs:
-            if not self._is_ext_call(e):
+            if not self._is_batched_fn_call(e):
                 continue
-            if current_ext_function is None or current_ext_function != e.fn:
+            if current_batched_fn is None or current_batched_fn != e.fn:
                 # create a new cohort
                 cohorts.append([])
-                current_ext_function = e.fn
+                current_batched_fn = e.fn
             cohorts[-1].append(e)
 
         # expand the cohorts to include all exprs that are in the same evaluation context as the external calls;
@@ -115,18 +115,18 @@ class ExprEvalNode(ExecNode):
             assert len(cohort) > 0
             # create the first segment here, so we can avoid checking for an empty list in the loop
             segments = [[cohort[0]]]
-            is_ext_segment = self._is_ext_call(cohort[0])
-            ext_fn: Optional[func.BatchedFunction] = self._get_batched_fn(cohort[0])
+            is_batched_segment = self._is_batched_fn_call(cohort[0])
+            batched_fn: Optional[CallableFunction] = self._get_batched_fn(cohort[0])
             for e in cohort[1:]:
-                if self._is_ext_call(e):
+                if self._is_batched_fn_call(e):
                     segments.append([e])
-                    is_ext_segment = True
-                    ext_fn = self._get_batched_fn(e)
+                    is_batched_segment = True
+                    batched_fn = self._get_batched_fn(e)
                 else:
-                    if is_ext_segment:
+                    if is_batched_segment:
                         # start a new segment
                         segments.append([])
-                        is_ext_segment = False
+                        is_batched_segment = False
                     segments[-1].append(e)
 
             # we create the EvalCtxs manually because create_eval_ctx() would repeat the dependencies of each segment
@@ -135,21 +135,21 @@ class ExprEvalNode(ExecNode):
                     slot_idxs=[e.slot_idx for e in s], exprs=s, target_slot_idxs=[], target_exprs=[])
                 for s in segments
             ]
-            cohort_info = self.Cohort(cohort, ext_fn, segment_ctxs, target_slot_idxs[i])
+            cohort_info = self.Cohort(cohort, batched_fn, segment_ctxs, target_slot_idxs[i])
             self.cohorts.append(cohort_info)
 
     def _exec_cohort(self, cohort: Cohort, rows: DataRowBatch) -> None:
         """Compute the cohort for the entire input batch by dividing it up into sub-batches"""
         batch_start_idx = 0  # start row of the current sub-batch
         # for multi-resolution models, we re-assess the correct ext fn batch size for each input batch
-        ext_batch_size = cohort.ext_function.get_batch_size() if cohort.ext_function is not None else None
+        ext_batch_size = cohort.batched_fn.get_batch_size() if cohort.batched_fn is not None else None
         if ext_batch_size is not None:
             cohort.batch_size = ext_batch_size
 
         while batch_start_idx < len(rows):
             num_batch_rows = min(cohort.batch_size, len(rows) - batch_start_idx)
             for segment_ctx in cohort.segment_ctxs:
-                if not self._is_ext_call(segment_ctx.exprs[0]):
+                if not self._is_batched_fn_call(segment_ctx.exprs[0]):
                     # compute batch row-wise
                     for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
                         self.row_builder.eval(
@@ -193,7 +193,7 @@ class ExprEvalNode(ExecNode):
                             for k in kwarg_batches.keys()
                         }
                         start_ts = time.perf_counter()
-                        result_batch = fn_call.fn.invoke(call_args, call_kwargs)
+                        result_batch = fn_call.fn.exec_batch(*call_args, **call_kwargs)
                         self.ctx.profile.eval_time[fn_call.slot_idx] += time.perf_counter() - start_ts
                         self.ctx.profile.eval_count[fn_call.slot_idx] += num_ext_batch_rows
 

pixeltable/exec/sql_scan_node.py

@@ -21,7 +21,6 @@ class SqlScanNode(ExecNode):
             select_list: Iterable[exprs.Expr],
             where_clause: Optional[exprs.Expr] = None, filter: Optional[exprs.Predicate] = None,
             order_by_items: Optional[List[Tuple[exprs.Expr, bool]]] = None,
-            similarity_clause: Optional[exprs.ImageSimilarityPredicate] = None,
             limit: int = 0, set_pk: bool = False, exact_version_only: Optional[List[catalog.TableVersion]] = None
     ):
         """
@@ -77,15 +76,17 @@ class SqlScanNode(ExecNode):
         # the number of tables that need to be joined to the target table
         for rowid_ref in [e for e, _ in order_by_items if isinstance(e, exprs.RowidRef)]:
             rowid_ref.set_tbl(tbl)
-        order_by_clause = [e.sql_expr().desc() if not asc else e.sql_expr() for e, asc in order_by_items]
+        order_by_clause: List[sql.ClauseElement] = []
+        for e, asc in order_by_items:
+            if isinstance(e, exprs.SimilarityExpr):
+                order_by_clause.append(e.as_order_by_clause(asc))
+            else:
+                order_by_clause.append(e.sql_expr().desc() if not asc else e.sql_expr())
 
         if where_clause is not None:
             sql_where_clause = where_clause.sql_expr()
             assert sql_where_clause is not None
             self.stmt = self.stmt.where(sql_where_clause)
-        if similarity_clause is not None:
-            self.stmt = self.stmt.order_by(
-                similarity_clause.img_col_ref.col.sa_idx_col.l2_distance(similarity_clause.embedding()))
         if len(order_by_clause) > 0:
             self.stmt = self.stmt.order_by(*order_by_clause)
         elif target.id in row_builder.unstored_iter_args:
@@ -201,7 +202,7 @@ class SqlScanNode(ExecNode):
                 self.row_builder.eval(output_row, self.filter_eval_ctx, profile=self.ctx.profile)
                 if output_row[self.filter.slot_idx]:
                     needs_row = True
-                    if self.limit is not None and len(output_batch) >= self.limit:
+                    if self.limit > 0 and len(output_batch) >= self.limit:
                         self.has_more_rows = False
                         break
                 else: