pixeltable 0.0.0__py3-none-any.whl

pixeltable/utils/__init__.py

@@ -0,0 +1,17 @@
+def print_perf_counter_delta(delta: float) -> str:
+    """Prints a performance counter delta in a human-readable format.
+
+    Args:
+        delta: delta in seconds
+
+    Returns:
+        Human-readable string
+    """
+    if delta < 1e-6:
+        return f'{delta * 1e9:.2f} ns'
+    elif delta < 1e-3:
+        return f'{delta * 1e6:.2f} us'
+    elif delta < 1:
+        return f'{delta * 1e3:.2f} ms'
+    else:
+        return f'{delta:.2f} s'

pixeltable/utils/arrow.py

@@ -0,0 +1,98 @@
+import logging
+from typing import Any, Dict, Iterable, Iterator, Optional
+
+import pyarrow as pa
+
+import pixeltable.type_system as ts
+
+_logger = logging.getLogger(__name__)
+
+_pa_to_pt: Dict[pa.DataType, ts.ColumnType] = {
+    pa.string(): ts.StringType(nullable=True),
+    pa.timestamp('us'): ts.TimestampType(nullable=True),
+    pa.bool_(): ts.BoolType(nullable=True),
+    pa.uint8(): ts.IntType(nullable=True),
+    pa.int8(): ts.IntType(nullable=True),
+    pa.uint32(): ts.IntType(nullable=True),
+    pa.uint64(): ts.IntType(nullable=True),
+    pa.int32(): ts.IntType(nullable=True),
+    pa.int64(): ts.IntType(nullable=True),
+    pa.float32(): ts.FloatType(nullable=True),
+}
+
+_pt_to_pa: Dict[ts.ColumnType, pa.DataType] = {
+    ts.StringType: pa.string(),
+    ts.TimestampType: pa.timestamp('us'),  # postgres timestamp is microseconds
+    ts.BoolType: pa.bool_(),
+    ts.IntType: pa.int64(),
+    ts.FloatType: pa.float32(),
+    ts.JsonType: pa.string(),  # TODO(orm) pa.struct() is possible
+    ts.ImageType: pa.binary(),  # inline image
+    ts.AudioType: pa.string(),  # path
+    ts.VideoType: pa.string(),  # path
+    ts.DocumentType: pa.string(),  # path
+}
+
+
+def to_pixeltable_type(arrow_type: pa.DataType) -> Optional[ts.ColumnType]:
+    """Convert a pyarrow DataType to a pixeltable ColumnType if one is defined.
+    Returns None if no conversion is currently implemented.
+    """
+    if arrow_type in _pa_to_pt:
+        return _pa_to_pt[arrow_type]
+    elif isinstance(arrow_type, pa.FixedShapeTensorType):
+        dtype = to_pixeltable_type(arrow_type.value_type)
+        if dtype is None:
+            return None
+        return ts.ArrayType(shape=arrow_type.shape, dtype=dtype)
+    else:
+        return None
+
+
+def to_arrow_type(pixeltable_type: ts.ColumnType) -> Optional[pa.DataType]:
+    """Convert a pixeltable DataType to a pyarrow datatype if one is defined.
+    Returns None if no conversion is currently implemented.
+    """
+    if pixeltable_type.__class__ in _pt_to_pa:
+        return _pt_to_pa[pixeltable_type.__class__]
+    elif isinstance(pixeltable_type, ts.ArrayType):
+        return pa.fixed_shape_tensor(pa.from_numpy_dtype(pixeltable_type.numpy_dtype()), pixeltable_type.shape)
+    else:
+        return None
+
+
+def to_pixeltable_schema(arrow_schema: pa.Schema) -> Dict[str, ts.ColumnType]:
+    return {field.name: to_pixeltable_type(field.type) for field in arrow_schema}
+
+
+def to_arrow_schema(pixeltable_schema: Dict[str, Any]) -> pa.Schema:
+    return pa.schema((name, to_arrow_type(typ)) for name, typ in pixeltable_schema.items())
+
+
+def to_pydict(batch: pa.RecordBatch) -> Dict[str, Iterable[Any]]:
+    """Convert a RecordBatch to a dictionary of lists, unlike pa.lib.RecordBatch.to_pydict,
+    this function will not convert numpy arrays to lists, and will preserve the original numpy dtype.
+    """
+    out = {}
+    for k, name in enumerate(batch.schema.names):
+        col = batch.column(k)
+        if isinstance(col.type, pa.FixedShapeTensorType):
+            # treat array columns as numpy arrays to easily preserve numpy type
+            out[name] = col.to_numpy(zero_copy_only=False)
+        else:
+            # for the rest, use pydict to preserve python types
+            out[name] = col.to_pylist()
+
+    return out
+
+
+def iter_tuples(batch: pa.RecordBatch) -> Iterator[Dict[str, Any]]:
+    """Convert a RecordBatch to an iterator of dictionaries. also works with pa.Table and pa.RowGroup"""
+    pydict = to_pydict(batch)
+    assert len(pydict) > 0, 'empty record batch'
+    for _, v in pydict.items():
+        batch_size = len(v)
+        break
+
+    for i in range(batch_size):
+        yield {col_name: values[i] for col_name, values in pydict.items()}

pixeltable/utils/clip.py

@@ -0,0 +1,18 @@
+import numpy as np
+import PIL.Image
+
+import pixeltable.func as func
+from pixeltable.env import Env
+
+
+def embed_image(img: PIL.Image.Image) -> np.ndarray:
+    from pixeltable.functions.nos.image_embedding import openai_clip
+    model_info = openai_clip.model_spec
+    result = Env.get().nos_client.Run(task=model_info.task, model_name=model_info.name, images=[img.resize((224, 224))])
+    return result['embedding'].squeeze(0)
+
+def embed_text(text: str) -> np.ndarray:
+    from pixeltable.functions.nos.text_embedding import openai_clip
+    model_info = openai_clip.model_spec
+    result = Env.get().nos_client.Run(task=model_info.task, model_name=model_info.name, texts=[text])
+    return result['embedding'].squeeze(0)

pixeltable/utils/coco.py

@@ -0,0 +1,136 @@
+from typing import List, Dict, Any, Set
+from pathlib import Path
+import json
+
+import PIL
+
+import pixeltable.exceptions as excs
+
+
+format_msg = """
+
+Required format:
+{
+    'image': PIL.Image.Image,
+    'annotations': [
+        {
+            'bbox': [x: int, y: int, w: int, h: int],
+            'category': str | int,
+        },
+        ...
+    ],
+}
+"""
+
+def _verify_input_dict(input_dict: Dict[str, Any]) -> None:
+    """Verify that input_dict is a valid input dict for write_coco_dataset()"""
+    if not isinstance(input_dict, dict):
+        raise excs.Error(f'Expected dict, got {input_dict}{format_msg}')
+    if 'image' not in input_dict:
+        raise excs.Error(f'Missing key "image" in input dict: {input_dict}{format_msg}')
+    if not isinstance(input_dict['image'], PIL.Image.Image):
+        raise excs.Error(f'Value for "image" is not a PIL.Image.Image: {input_dict}{format_msg}')
+    if  'annotations' not in input_dict:
+        raise excs.Error(f'Missing key "annotations" in input dict: {input_dict}{format_msg}')
+    if not isinstance(input_dict['annotations'], list):
+        raise excs.Error(f'Value for "annotations" is not a list: {input_dict}{format_msg}')
+    for annotation in input_dict['annotations']:
+        if not isinstance(annotation, dict):
+            raise excs.Error(f'Annotation is not a dict: {annotation}{format_msg}')
+        if 'bbox' not in annotation:
+            raise excs.Error(f'Missing key "bbox" in annotation: {annotation}{format_msg}')
+        if not isinstance(annotation['bbox'], list):
+            raise excs.Error(f'Value for "bbox" is not a list [x, y, w, h]: {annotation}{format_msg}')
+        if len(annotation['bbox']) != 4 or not all(isinstance(x, int) for x in annotation['bbox']):
+            raise excs.Error(f'Key "bbox" is not a list [x, y, w, h] of ints: {annotation}{format_msg}')
+        if 'category' not in annotation:
+            raise excs.Error(f'Missing key "category" in annotation: {annotation}{format_msg}')
+        if not isinstance(annotation['category'], (str, int)):
+            raise excs.Error(f'Value for "category" is not a str or int: {annotation}{format_msg}')
+
+def write_coco_dataset(df: 'pixeltable.DataFrame', dest_path: Path) -> Path:
+    """Export a DataFrame result set as a COCO dataset in dest_path and return the path of the data.json file."""
+    # TODO: validate schema
+    if len(df._select_list_exprs) != 1 or not df._select_list_exprs[0].col_type.is_json_type():
+        raise excs.Error(f'Expected exactly one json-typed column in select list: {df._select_list_exprs}')
+    input_dict_slot_idx = -1  # df._select_list_exprs[0].slot_idx isn't valid until _exec()
+
+    # create output dir
+    assert not dest_path.exists()
+    dest_path.mkdir(parents=False)
+    images_dir = dest_path / 'images'
+    images_dir.mkdir()
+
+    images: List[Dict[str, Any]] = []
+    img_id = -1
+    annotations: List[Dict[str, Any]] = []
+    ann_id = -1
+    categories: Set[Any] = set()
+    for input_row in df._exec():
+        if input_dict_slot_idx == -1:
+            input_dict_expr = df._select_list_exprs[0]
+            input_dict_slot_idx = input_dict_expr.slot_idx
+            input_dict = input_row[input_dict_slot_idx]
+            _verify_input_dict(input_dict)
+
+            # we want to know the slot idx of the image used in the input dict, so that we can check whether we
+            # already have a local path for it
+            input_dict_dependencies = input_dict_expr.dependencies()
+            img_slot_idx = next((e.slot_idx for e in input_dict_dependencies if e.col_type.is_image_type()), None)
+            assert img_slot_idx is not None
+        else:
+            input_dict = input_row[input_dict_slot_idx]
+            _verify_input_dict(input_dict)
+
+        # create image record
+        img_id += 1
+
+        # get a local path for the image
+        img = input_dict['image']
+        if input_row.file_paths[img_slot_idx] is not None:
+            # we already have a local path
+            img_path = Path(input_row.file_paths[img_slot_idx])
+            # TODO: if the path leads to our tmp dir, we need to move the file
+        else:
+            # we need to create a local path
+            img_path = images_dir / f'{img_id}.jpg'
+            img.save(img_path)
+
+        images.append({
+            'id': img_id,
+            'file_name': str(img_path),
+            'width': img.width,
+            'height': img.height,
+        })
+
+        # create annotation records for this image
+        for annotation in input_dict['annotations']:
+            ann_id += 1
+            x, y, w, h = annotation['bbox']
+            category = annotation['category']
+            categories.add(category)
+            annotations.append({
+                'id': ann_id,
+                'image_id': img_id,
+                # we use the category name here and fix it up at the end, when we have assigned category ids
+                'category_id': category,
+                'bbox': annotation['bbox'],
+                'area': w * h,
+                'iscrowd': 0,
+            })
+
+    # replace category names with ids
+    category_ids = {category: id for id, category in enumerate(sorted(list(categories)))}
+    for annotation in annotations:
+        annotation['category_id'] = category_ids[annotation['category_id']]
+
+    result = {
+        'images': images,
+        'annotations': annotations,
+        'categories': [{'id': id, 'name': category} for category, id in category_ids.items()],
+    }
+    output_path = dest_path / 'data.json'
+    with open(output_path, 'w') as f:
+        json.dump(result, f)
+    return output_path
+

pixeltable/utils/documents.py

@@ -0,0 +1,69 @@
+from typing import Optional, Dict
+import dataclasses
+
+import pixeltable.type_system as ts
+
+
+@dataclasses.dataclass
+class DocumentHandle:
+    format: ts.DocumentType.DocumentFormat
+    bs_doc: Optional['bs4.BeautifulSoup'] = None
+    md_ast: Optional[Dict] = None
+    pdf_doc: Optional['fitz.Document'] = None
+
+def get_document_handle(path: str) -> Optional[DocumentHandle]:
+    # try pdf first, because a correct PDF is a binary format that
+    # would trigger encoding exceptions if oppened as utf8.
+    pdf_doc = get_pdf_handle(path)
+    if pdf_doc is not None:
+        return DocumentHandle(format=ts.DocumentType.DocumentFormat.PDF, pdf_doc=pdf_doc)
+    # currently the rest of the types are text-based, so we can open them in utf8 mode once
+    try:
+        with open(path, 'r', encoding='utf8') as file:
+            contents = file.read()
+    except UnicodeDecodeError:
+        # not pdf, and also not valid text file
+        return None
+
+    # bs4 will appear to succeed for md files as well.
+    # this will break most markdown files at the moment.
+    bs_doc = get_html_handle(contents)
+    if bs_doc is not None:
+        return DocumentHandle(format=ts.DocumentType.DocumentFormat.HTML, bs_doc=bs_doc)
+
+    md_ast = get_markdown_handle(contents)
+    if md_ast is not None:
+        return DocumentHandle(format=ts.DocumentType.DocumentFormat.MD, md_ast=md_ast)
+
+    return None
+
+def get_html_handle(text: str) -> Optional['bs4.BeautifulSoup']:
+    import bs4
+    try:
+        doc = bs4.BeautifulSoup(text, 'html.parser')
+        if doc.find() is None:
+            return None
+        return doc
+    except Exception:
+        return None
+
+def get_markdown_handle(text: str) -> Optional[Dict]:
+    import mistune
+    try:
+        md_ast = mistune.create_markdown(renderer=None)
+        return md_ast(text)
+    except Exception:
+        return None
+
+def get_pdf_handle(path : str) -> Optional['fitz.Document']:
+    import fitz  # aka pymupdf
+    try:
+        doc = fitz.open(path)
+        # check pdf (bc it will work for images)
+        if not doc.is_pdf:
+            return None
+        # try to read one page
+        next(page for page in doc)
+        return doc
+    except Exception:
+        return None

pixeltable/utils/filecache.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+from typing import Optional, List, Tuple, Dict
+from collections import OrderedDict, defaultdict, namedtuple
+import os
+import glob
+from pathlib import Path
+from time import time
+import logging
+from uuid import UUID
+import hashlib
+
+from pixeltable.env import Env
+
+
+_logger = logging.getLogger('pixeltable')
+
+class CacheEntry:
+    def __init__(self, key: str, tbl_id: UUID, col_id: int, size: int, last_accessed_ts: int, ext: str):
+        self.key = key
+        self.tbl_id = tbl_id
+        self.col_id = col_id
+        self.size = size
+        self.last_accessed_ts = last_accessed_ts
+        self.ext = ext
+
+    def path(self) -> Path:
+        return Env.get().file_cache_dir / f'{self.tbl_id.hex}_{self.col_id}_{self.key}{self.ext}'
+
+    @classmethod
+    def from_file(cls, path: Path) -> CacheEntry:
+        components = path.stem.split('_')
+        assert len(components) == 3
+        tbl_id = UUID(components[0])
+        col_id = int(components[1])
+        key = components[2]
+        file_info = os.stat(str(path))
+        return cls(key, tbl_id, col_id, file_info.st_size, file_info.st_mtime, path.suffix)
+
+
+class FileCache:
+    """
+    A local cache of external (eg, S3) file references in cells of a stored table (ie, table or view).
+
+    Cache entries are identified by a hash of the file url and stored in Env.filecache_dir. The time of last
+    access of a cache entries is its file's mtime.
+
+    TODO:
+    - enforce a maximum capacity with LRU eviction
+    - implement MRU eviction for queries that exceed the capacity
+    """
+    _instance: Optional[FileCache] = None
+    ColumnStats = namedtuple('FileCacheColumnStats', ['tbl_id', 'col_id', 'num_files', 'total_size'])
+    CacheStats = namedtuple(
+        'FileCacheStats', ['total_size', 'num_requests', 'num_hits', 'num_evictions', 'column_stats'])
+
+    @classmethod
+    def get(cls) -> FileCache:
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        self.cache: OrderedDict[str, CacheEntry] = OrderedDict()  # ordered by entry.last_accessed_ts
+        self.total_size = 0
+        #self.capacity = Env.get().max_filecache_size
+        self.num_requests = 0
+        self.num_hits = 0
+        self.num_evictions = 0
+        paths = glob.glob(str(Env.get().file_cache_dir / '*'))
+        entries = [CacheEntry.from_file(Path(path_str)) for path_str in paths]
+        # we need to insert entries in order of last_accessed_ts
+        entries.sort(key=lambda e: e.last_accessed_ts)
+        for entry in entries:
+            self.cache[entry.key] = entry
+            self.total_size += entry.size
+
+    def avg_file_size(self) -> int:
+        if len(self.cache) == 0:
+            return 0
+        return int(self.total_size / len(self.cache))
+
+    def num_files(self, tbl_id: Optional[UUID] = None) -> int:
+        if tbl_id is None:
+            return len(self.cache)
+        entries = [e for e in self.cache.values() if e.tbl_id == tbl_id]
+        return len(entries)
+
+    def clear(self, tbl_id: Optional[UUID] = None, capacity: Optional[int] = None) -> None:
+        """
+        For testing purposes: allow resetting capacity and stats.
+        """
+        self.num_requests, self.num_hits, self.num_evictions = 0, 0, 0
+        entries = list(self.cache.values())  # list(): avoid dealing with values() return type
+        if tbl_id is not None:
+            entries = [e for e in entries if e.tbl_id == tbl_id]
+            _logger.debug(f'clearing {len(entries)} entries from file cache for table {tbl_id}')
+        else:
+            _logger.debug(f'clearing {len(entries)} entries from file cache')
+        for entry in entries:
+            del self.cache[entry.key]
+            self.total_size -= entry.size
+            os.remove(entry.path())
+        # if capacity is not None:
+        #     self.capacity = capacity
+        # else:
+        #     # need to reset to default
+        #     self.capacity = Env.get().max_filecache_size
+        # _logger.debug(f'setting file cache capacity to {self.capacity}')
+
+    def _url_hash(self, url: str) -> str:
+        h = hashlib.sha256()
+        h.update(url.encode())
+        return h.hexdigest()
+
+    def lookup(self, url: str) -> Optional[Path]:
+        self.num_requests += 1
+        key = self._url_hash(url)
+        entry = self.cache.get(key, None)
+        if entry is None:
+            _logger.debug(f'file cache miss for {url}')
+            return None
+        # update mtime and cache
+        path = entry.path()
+        path.touch(exist_ok=True)
+        file_info = os.stat(str(path))
+        entry.last_accessed_ts = file_info.st_mtime
+        self.cache.move_to_end(key, last=True)
+        self.num_hits += 1
+        _logger.debug(f'file cache hit for {url}')
+        return path
+
+    # def can_admit(self, query_ts: int) -> bool:
+    #     if self.total_size + self.avg_file_size <= self.capacity:
+    #         return True
+    #     assert len(self.cache) > 0
+    #     # check whether we can evict the current lru entry
+    #     lru_entry = next(iter(self.cache.values()))
+    #     if lru_entry.last_accessed_ts >= query_ts:
+    #         # the current query brought this entry in: we're not going to evict it
+    #         return False
+    #     return True
+
+    def add(self, tbl_id: UUID, col_id: int, url: str, path: Path) -> Path:
+        """Adds url at 'path' to cache and returns its new path.
+        'path' will not be accessible after this call. Retains the extension of 'path'.
+        """
+        file_info = os.stat(str(path))
+        _ = time()
+        #if self.total_size + file_info.st_size > self.capacity:
+        if False:
+            if len(self.cache) == 0:
+                # nothing to evict
+                return
+            # evict entries until we're below the limit or until we run into entries the current query brought in
+            while True:
+                lru_entry = next(iter(self.cache.values()))
+                if lru_entry.last_accessed_ts >= query_ts:
+                    # the current query brought this entry in: switch to MRU and ignore this put()
+                    _logger.debug('file cache switched to MRU')
+                    return
+                self.cache.popitem(last=False)
+                self.total_size -= lru_entry.size
+                self.num_evictions += 1
+                os.remove(str(lru_entry.path()))
+                _logger.debug(f'evicted entry for cell {lru_entry.cell_id} from file cache')
+                if self.total_size + file_info.st_size <= self.capacity:
+                    break
+
+        key = self._url_hash(url)
+        assert key not in self.cache
+        entry = CacheEntry(key, tbl_id, col_id, file_info.st_size, file_info.st_mtime, path.suffix)
+        self.cache[key] = entry
+        self.total_size += entry.size
+        new_path = entry.path()
+        os.rename(str(path), str(new_path))
+        _logger.debug(f'added entry for cell {url} to file cache')
+        return new_path
+
+    def stats(self) -> CacheStats:
+        # collect column stats
+        # (tbl_id, col_id) -> (num_files, total_size)
+        d: Dict[Tuple[int, int], List[int]] = defaultdict(lambda: [0, 0])
+        for entry in self.cache.values():
+            t = d[(entry.tbl_id, entry.col_id)]
+            t[0] += 1
+            t[1] += entry.size
+        col_stats = [
+            self.ColumnStats(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()
+        ]
+        col_stats.sort(key=lambda e: e[3], reverse=True)
+        return self.CacheStats(self.total_size, self.num_requests, self.num_hits, self.num_evictions, col_stats)
+
+    def debug_print(self) -> None:
+        for entry in self.cache.values():
+            print(f'CacheEntry: tbl_id={entry.tbl_id}, col_id={entry.col_id}, size={entry.size}')

pixeltable/utils/help.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+import pixeltable.func as func
+
+
+def help(obj: Any) -> None:
+    """Returns help text for the given object."""
+    if isinstance(obj, func.Function):
+        print(obj.help_str())
+    else:
+        print(__builtins__.help(obj))

pixeltable/utils/http_server.py

@@ -0,0 +1,70 @@
+import http
+import http.server
+import logging
+import urllib
+import posixpath
+import pathlib
+import os
+import string
+
+_logger = logging.getLogger('pixeltable.http.server')
+
+
+def get_file_uri(http_address: str, file_path: str) -> str:
+    """Get the URI for a file path, with the given prefix.
+        Used in the client to generate a URI
+    """
+    abs_path = pathlib.Path(file_path)
+    assert abs_path.is_absolute()
+    url = urllib.request.pathname2url(str(abs_path))
+    return f'{http_address}{url}'
+
+
+class AbsolutePathHandler(http.server.SimpleHTTPRequestHandler):
+    """Serves all absolute paths, not just the current directory"""
+    def translate_path(self, path: str) -> str:
+        """
+            Translate a /-separated PATH to the local filename syntax.
+            overrides http.server.SimpleHTTPRequestHandler.translate_path
+
+            This is only useful for file serving.
+
+            Code initially taken from there:
+            https://github.com/python/cpython/blob/f5406ef454662b98df107775d18ff71ae6849618/Lib/http/server.py#L834
+        """
+        _logger.info(f'translate path {path=}')
+        # abandon query parameters, taken from http.server.SimpleHTTPRequestHandler
+        path = path.split('?', 1)[0]
+        path = path.split('#', 1)[0]
+
+        path = pathlib.Path(urllib.request.url2pathname(path))
+        return str(path)
+
+    def log_message(self, format, *args) -> None:
+        """override logging to stderr in http.server.BaseHTTPRequestHandler"""
+        message = format % args
+        _logger.info(message.translate(self._control_char_table))
+
+
+class LoggingHTTPServer(http.server.ThreadingHTTPServer):
+    """Avoids polluting stdout and stderr"""
+
+    def handle_error(self, request, client_address) -> None:
+        """override socketserver.TCPServer.handle_error which prints directly to sys.stderr"""
+        import traceback
+
+        _logger.error(
+            f'Exception occurred during processing of {request=} from {client_address=}\
+            \nbacktrace:\n{traceback.format_exc()}\n----\n'
+        )
+
+
+def make_server(address: str, port: int) -> http.server.HTTPServer:
+    """Create a file server with pixeltable specific config """
+    return LoggingHTTPServer((address, port), AbsolutePathHandler)
+
+
+if __name__ == '__main__':
+    httpd = make_server('127.0.0.1', 8000)
+    print(f'about to server HTTP on {httpd.server_address}')
+    httpd.serve_forever()