pixeltable 0.4.3__py3-none-any.whl → 0.4.5__py3-none-any.whl

pixeltable/exec/data_row_batch.py

@@ -3,8 +3,7 @@ from __future__ import annotations
 import logging
 from typing import Iterator, Optional
 
-from pixeltable import catalog, exprs
-from pixeltable.utils.media_store import MediaStore
+from pixeltable import exprs
 
 _logger = logging.getLogger('pixeltable')
 
@@ -15,51 +14,19 @@ class DataRowBatch:
     Contains the metadata needed to initialize DataRows.
     """
 
-    tbl: Optional[catalog.TableVersionHandle]
     row_builder: exprs.RowBuilder
-    img_slot_idxs: list[int]
-    media_slot_idxs: list[int]  # non-image media slots
-    array_slot_idxs: list[int]
     rows: list[exprs.DataRow]
 
-    def __init__(
-        self,
-        tbl: Optional[catalog.TableVersionHandle],
-        row_builder: exprs.RowBuilder,
-        num_rows: Optional[int] = None,
-        rows: Optional[list[exprs.DataRow]] = None,
-    ):
+    def __init__(self, row_builder: exprs.RowBuilder, rows: Optional[list[exprs.DataRow]] = None):
         """
         Requires either num_rows or rows to be specified, but not both.
         """
-        assert num_rows is None or rows is None
-        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
-        self.media_slot_idxs = [
-            e.slot_idx
-            for e in row_builder.unique_exprs
-            if e.col_type.is_media_type() and not e.col_type.is_image_type()
-        ]
-        self.array_slot_idxs = [e.slot_idx for e in row_builder.unique_exprs if e.col_type.is_array_type()]
-        if rows is not None:
-            self.rows = rows
-        else:
-            if num_rows is None:
-                num_rows = 0
-            self.rows = [
-                exprs.DataRow(
-                    row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
-                )
-                for _ in range(num_rows)
-            ]
+        self.rows = [] if rows is None else rows
 
-    def add_row(self, row: Optional[exprs.DataRow] = None) -> exprs.DataRow:
+    def add_row(self, row: Optional[exprs.DataRow]) -> exprs.DataRow:
         if row is None:
-            row = exprs.DataRow(
-                self.row_builder.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs
-            )
+            row = self.row_builder.make_row()
         self.rows.append(row)
         return row
 
@@ -73,26 +40,18 @@ class DataRowBatch:
         return self.rows[index]
 
     def flush_imgs(
-        self,
-        idx_range: Optional[slice] = None,
-        stored_img_info: Optional[list[exprs.ColumnSlotIdx]] = None,
-        flushed_slot_idxs: Optional[list[int]] = None,
+        self, idx_range: Optional[slice], stored_img_info: list[exprs.ColumnSlotIdx], flushed_img_slots: list[int]
     ) -> 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:
-            flushed_slot_idxs = []
-        if len(stored_img_info) == 0 and len(flushed_slot_idxs) == 0:
+        if len(stored_img_info) == 0 and len(flushed_img_slots) == 0:
             return
+
         if idx_range is None:
             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.get().version))
-                row.flush_img(info.slot_idx, filepath)
-            for slot_idx in flushed_slot_idxs:
+                row.flush_img(info.slot_idx, info.col)
+            for slot_idx in flushed_img_slots:
                 row.flush_img(slot_idx)
 
     def __iter__(self) -> Iterator[exprs.DataRow]:

pixeltable/exec/expr_eval/expr_eval_node.py

@@ -240,7 +240,7 @@ class ExprEvalNode(ExecNode):
                         # make sure we top up our in-flight rows before yielding
                         self._dispatch_input_rows()
                         self._log_state(f'yielding {len(batch_rows)} rows')
-                        yield DataRowBatch(tbl=None, row_builder=self.row_builder, rows=batch_rows)
+                        yield DataRowBatch(row_builder=self.row_builder, rows=batch_rows)
                         # at this point, we may have more completed rows
 
                 assert self.completed_rows.empty()  # all completed rows should be sitting in output_buffer
@@ -254,7 +254,7 @@ class ExprEvalNode(ExecNode):
                         batch_rows = self.output_buffer.get_rows(self.output_buffer.num_ready)
                         self.num_output_rows += len(batch_rows)
                         self._log_state(f'yielding {len(batch_rows)} rows')
-                        yield DataRowBatch(tbl=None, row_builder=self.row_builder, rows=batch_rows)
+                        yield DataRowBatch(row_builder=self.row_builder, rows=batch_rows)
 
                     assert self.output_buffer.num_rows == 0
                     return

pixeltable/exec/in_memory_data_node.py

@@ -23,7 +23,7 @@ class InMemoryDataNode(ExecNode):
 
     input_rows: list[dict[str, Any]]
     start_row_id: int
-    output_rows: Optional[DataRowBatch]
+    output_batch: Optional[DataRowBatch]
 
     # output_exprs is declared in the superclass, but we redeclare it here with a more specific type
     output_exprs: list[exprs.ColumnRef]
@@ -42,7 +42,7 @@ class InMemoryDataNode(ExecNode):
         self.tbl = tbl
         self.input_rows = rows
         self.start_row_id = start_row_id
-        self.output_rows = None
+        self.output_batch = None
 
     def _open(self) -> None:
         """Create row batch and populate with self.input_rows"""
@@ -56,22 +56,22 @@ class InMemoryDataNode(ExecNode):
         }
         output_slot_idxs = {e.slot_idx for e in self.output_exprs}
 
-        self.output_rows = DataRowBatch(self.tbl, self.row_builder, len(self.input_rows))
-        for row_idx, input_row in enumerate(self.input_rows):
+        self.output_batch = DataRowBatch(self.row_builder)
+        for input_row in self.input_rows:
+            output_row = self.row_builder.make_row()
             # populate the output row with the values provided in the input row
             input_slot_idxs: set[int] = set()
             for col_name, val in input_row.items():
                 col_info = user_cols_by_name.get(col_name)
                 assert col_info is not None
-
-                if col_info.col.col_type.is_image_type() and isinstance(val, bytes):
-                    # this is a literal image, ie, a sequence of bytes; we save this as a media file and store the path
-                    path = str(MediaStore.prepare_media_path(self.tbl.id, col_info.col.id, self.tbl.get().version))
-                    with open(path, 'wb') as fp:
-                        fp.write(val)
-                    self.output_rows[row_idx][col_info.slot_idx] = path
+                col = col_info.col
+                if col.col_type.is_image_type() and isinstance(val, bytes):
+                    # this is a literal media file, ie, a sequence of bytes; save it as a binary file and store the path
+                    assert col.tbl.id == self.tbl.id
+                    filepath, _ = MediaStore.save_media_object(val, col, format=None)
+                    output_row[col_info.slot_idx] = str(filepath)
                 else:
-                    self.output_rows[row_idx][col_info.slot_idx] = val
+                    output_row[col_info.slot_idx] = val
 
                 input_slot_idxs.add(col_info.slot_idx)
 
@@ -80,10 +80,11 @@ class InMemoryDataNode(ExecNode):
             for slot_idx in missing_slot_idxs:
                 col_info = output_cols_by_idx.get(slot_idx)
                 assert col_info is not None
-                self.output_rows[row_idx][col_info.slot_idx] = None
+                output_row[col_info.slot_idx] = None
+            self.output_batch.add_row(output_row)
 
-        self.ctx.num_rows = len(self.output_rows)
+        self.ctx.num_rows = len(self.output_batch)
 
     async def __aiter__(self) -> AsyncIterator[DataRowBatch]:
-        _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_rows)} output_rows')
-        yield self.output_rows
+        _logger.debug(f'InMemoryDataNode: created row batch with {len(self.output_batch)} rows')
+        yield self.output_batch

pixeltable/exec/sql_node.py

@@ -316,8 +316,7 @@ class SqlNode(ExecNode):
             for _ in w:
                 pass
 
-        tbl_version = self.tbl.tbl_version if self.tbl is not None else None
-        output_batch = DataRowBatch(tbl_version, self.row_builder)
+        output_batch = DataRowBatch(self.row_builder)
         output_row: Optional[exprs.DataRow] = None
         num_rows_returned = 0
 
@@ -359,7 +358,7 @@ class SqlNode(ExecNode):
             if self.ctx.batch_size > 0 and len(output_batch) == self.ctx.batch_size:
                 _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
                 yield output_batch
-                output_batch = DataRowBatch(tbl_version, self.row_builder)
+                output_batch = DataRowBatch(self.row_builder)
 
         if len(output_batch) > 0:
             _logger.debug(f'SqlScanNode: returning {len(output_batch)} rows')
@@ -569,10 +568,10 @@ class SqlSampleNode(SqlNode):
         General SQL form is:
         - MD5(<seed::text> [ + '___' + <rowid_col_val>::text]+
         """
-        sql_expr: sql.ColumnElement = sql.cast(seed, sql.Text)
+        sql_expr: sql.ColumnElement = seed.cast(sql.String)
         for e in sql_cols:
             # Quotes are required below to guarantee that the string is properly presented in SQL
-            sql_expr = sql_expr + sql.literal_column("'___'", sql.Text) + sql.cast(e, sql.Text)
+            sql_expr = sql_expr + sql.literal_column("'___'", sql.Text) + e.cast(sql.String)
         sql_expr = sql.func.md5(sql_expr)
         return sql_expr
 
@@ -591,9 +590,9 @@ class SqlSampleNode(SqlNode):
                 s_key = self._create_key_sql(self.input_cte)
 
                 # Construct a suitable where clause
-                fraction_sql = sql.cast(SampleClause.fraction_to_md5_hex(float(self.sample_clause.fraction)), sql.Text)
+                fraction_md5 = SampleClause.fraction_to_md5_hex(self.sample_clause.fraction)
                 order_by = self._create_key_sql(self.input_cte)
-                return sql.select(*self.input_cte.c).where(s_key < fraction_sql).order_by(order_by)
+                return sql.select(*self.input_cte.c).where(s_key < fraction_md5).order_by(order_by)
 
             return self._create_stmt_stratified_fraction(self.sample_clause.fraction)
         else:

pixeltable/exprs/column_ref.py

@@ -325,7 +325,8 @@ class ColumnRef(Expr):
     @classmethod
     def get_column(cls, d: dict) -> catalog.Column:
         tbl_id, version, col_id = UUID(d['tbl_id']), d['tbl_version'], d['col_id']
-        tbl_version = catalog.Catalog.get().get_tbl_version(tbl_id, version)
+        # validate_initialized=False: this gets called as part of TableVersion.init()
+        tbl_version = catalog.Catalog.get().get_tbl_version(tbl_id, version, validate_initialized=False)
         # don't use tbl_version.cols_by_id here, this might be a snapshot reference to a column that was then dropped
         col = next(col for col in tbl_version.cols if col.id == col_id)
         return col

pixeltable/exprs/data_row.py

@@ -13,7 +13,8 @@ import PIL
 import PIL.Image
 import sqlalchemy as sql
 
-from pixeltable import env
+from pixeltable import catalog, env
+from pixeltable.utils.media_store import MediaStore
 
 
 class DataRow:
@@ -256,23 +257,22 @@ class DataRow:
             self.vals[idx] = val
         self.has_val[idx] = True
 
-    def flush_img(self, index: int, filepath: Optional[str] = None) -> None:
-        """Discard the in-memory value and save it to a local file, if filepath is not None"""
+    def flush_img(self, index: int, col: Optional[catalog.Column] = None) -> None:
+        """Save or discard the in-memory value (required to be a PIL.Image.Image)"""
         if self.vals[index] is None:
             return
         assert self.excs[index] is None
         if self.file_paths[index] is None:
-            if filepath is not None:
+            if col is not None:
                 image = self.vals[index]
-                assert isinstance(image, PIL.Image.Image)
-                # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
-                # In that case, use WebP instead.
-                format = 'webp' if image.has_transparency_data else 'jpeg'
-                if not filepath.endswith(f'.{format}'):
-                    filepath += f'.{format}'
-                self.file_paths[index] = filepath
-                self.file_urls[index] = urllib.parse.urljoin('file:', urllib.request.pathname2url(filepath))
-                image.save(filepath, format=format)
+                format = None
+                if isinstance(image, PIL.Image.Image):
+                    # Default to JPEG unless the image has a transparency layer (which isn't supported by JPEG).
+                    # In that case, use WebP instead.
+                    format = 'webp' if image.has_transparency_data else 'jpeg'
+                filepath, url = MediaStore.save_media_object(image, col, format=format)
+                self.file_paths[index] = str(filepath)
+                self.file_urls[index] = url
             else:
                 # we discard the content of this cell
                 self.has_val[index] = False

pixeltable/exprs/row_builder.py

@@ -8,9 +8,8 @@ from uuid import UUID
 
 import numpy as np
 
-from pixeltable import catalog, exceptions as excs, utils
+from pixeltable import catalog, exceptions as excs, exprs, utils
 from pixeltable.env import Env
-from pixeltable.utils.media_store import MediaStore
 
 from .data_row import DataRow
 from .expr import Expr, ExprScope
@@ -85,6 +84,10 @@ class RowBuilder:
     # (a subexpr can be shared across multiple output exprs)
     output_expr_ids: list[set[int]]
 
+    img_slot_idxs: list[int]  # Indices of image slots
+    media_slot_idxs: list[int]  # Indices of non-image media slots
+    array_slot_idxs: list[int]  # Indices of array slots
+
     @dataclass
     class EvalCtx:
         """Context for evaluating a set of target exprs"""
@@ -235,6 +238,12 @@ class RowBuilder:
         for e in self.output_exprs:
             self._record_output_expr_id(e, e.slot_idx)
 
+        self.img_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_image_type()]
+        self.media_slot_idxs = [
+            e.slot_idx for e in self.unique_exprs if e.col_type.is_media_type() and not e.col_type.is_image_type()
+        ]
+        self.array_slot_idxs = [e.slot_idx for e in self.unique_exprs if e.col_type.is_array_type()]
+
     def add_table_column(self, col: catalog.Column, slot_idx: int) -> None:
         """Record a column that is part of the table row"""
         assert self.tbl is not None
@@ -462,8 +471,7 @@ class RowBuilder:
             else:
                 if col.col_type.is_image_type() and data_row.file_urls[slot_idx] is None:
                     # we have yet to store this image
-                    filepath = str(MediaStore.prepare_media_path(col.tbl.id, col.id, col.tbl.version))
-                    data_row.flush_img(slot_idx, filepath)
+                    data_row.flush_img(slot_idx, col)
                 val = data_row.get_stored_val(slot_idx, col.get_sa_col_type())
                 table_row.append(val)
                 if col.stores_cellmd:
@@ -489,3 +497,7 @@ class RowBuilder:
                 store_col_names.append(col.col.cellmd_store_name())
 
         return store_col_names, media_cols
+
+    def make_row(self) -> exprs.DataRow:
+        """Creates a new DataRow with the current row_builder's configuration."""
+        return exprs.DataRow(self.num_materialized, self.img_slot_idxs, self.media_slot_idxs, self.array_slot_idxs)

pixeltable/exprs/string_op.py

@@ -68,7 +68,7 @@ class StringOp(Expr):
         if self.operator == StringOperator.CONCAT:
             return left.concat(right)
         if self.operator == StringOperator.REPEAT:
-            return sql.func.repeat(sql.cast(left, sql.String), sql.cast(right, sql.Integer))
+            return sql.func.repeat(left.cast(sql.String), right.cast(sql.Integer))
         return None
 
     def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:

pixeltable/func/expr_template_function.py

@@ -101,13 +101,10 @@ class ExprTemplateFunction(Function):
         return None
 
     def exec(self, args: Sequence[Any], kwargs: dict[str, Any]) -> Any:
-        from pixeltable import exec
-
         assert not self.is_polymorphic
         expr = self.instantiate(args, kwargs)
         row_builder = exprs.RowBuilder(output_exprs=[expr], columns=[], input_exprs=[])
-        row_batch = exec.DataRowBatch(tbl=None, row_builder=row_builder, num_rows=1)
-        row = row_batch[0]
+        row = row_builder.make_row()
         row_builder.eval(row, ctx=row_builder.default_eval_ctx)
         return row[row_builder.get_output_exprs()[0].slot_idx]
 

pixeltable/functions/date.py

@@ -83,7 +83,7 @@ def make_date(year: int, month: int, day: int) -> date:
 
 @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))
+    return sql.func.make_date(year.cast(sql.Integer), month.cast(sql.Integer), day.cast(sql.Integer))
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/gemini.py

@@ -7,7 +7,6 @@ the [Working with Gemini](https://pixeltable.readme.io/docs/working-with-gemini)
 
 import asyncio
 import io
-import tempfile
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional
 
@@ -215,9 +214,10 @@ async def generate_videos(
     video_bytes = await _genai_client().aio.files.download(file=video.video)  # type: ignore[arg-type]
     assert video_bytes is not None
 
-    _, output_filename = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
-    Path(output_filename).write_bytes(video_bytes)
-    return output_filename
+    # Create a temporary file to store the video bytes
+    output_path = env.Env.get().create_tmp_path('.mp4')
+    Path(output_path).write_bytes(video_bytes)
+    return str(output_path)
 
 
 @generate_videos.resource_pool

pixeltable/functions/math.py

@@ -97,7 +97,7 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
     if digits is None:
         return sql.func.round(self)
     else:
-        return sql.func.round(sql.cast(self, sql.Numeric), sql.cast(digits, sql.Integer))
+        return sql.func.round(self.cast(sql.Numeric), digits.cast(sql.Integer))
 
 
 @pxt.udf(is_method=True)

pixeltable/functions/openai.py

@@ -13,7 +13,6 @@ import logging
 import math
 import pathlib
 import re
-import uuid
 from typing import TYPE_CHECKING, Any, Callable, Optional, Type
 
 import httpx
@@ -32,11 +31,15 @@ _logger = logging.getLogger('pixeltable')
 
 
 @env.register_client('openai')
-def _(api_key: str) -> 'openai.AsyncOpenAI':
+def _(api_key: str, base_url: Optional[str] = None, api_version: Optional[str] = None) -> 'openai.AsyncOpenAI':
     import openai
 
+    default_query = None if api_version is None else {'api-version': api_version}
+
     return openai.AsyncOpenAI(
         api_key=api_key,
+        base_url=base_url,
+        default_query=default_query,
         # recommended to increase limits for async client to avoid connection errors
         http_client=httpx.AsyncClient(limits=httpx.Limits(max_keepalive_connections=100, max_connections=500)),
     )
@@ -125,7 +128,7 @@ _header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d
 def _parse_header_duration(duration_str: str) -> datetime.timedelta:
     match = _header_duration_pattern.match(duration_str)
     if not match:
-        raise ValueError('Invalid duration format')
+        raise ValueError(f'Invalid duration format: {duration_str}')
 
     days = int(match.group(1) or 0)
     hours = int(match.group(2) or 0)
@@ -148,7 +151,7 @@ def _get_header_info(
         requests_limit = int(requests_limit_str) if requests_limit_str is not None else None
         requests_remaining_str = headers.get('x-ratelimit-remaining-requests')
         requests_remaining = int(requests_remaining_str) if requests_remaining_str is not None else None
-        requests_reset_str = headers.get('x-ratelimit-reset-requests')
+        requests_reset_str = headers.get('x-ratelimit-reset-requests', '5s')  # Default to 5 seconds
         requests_reset_ts = now + _parse_header_duration(requests_reset_str)
         requests_info = (requests_limit, requests_remaining, requests_reset_ts)
 
@@ -158,7 +161,7 @@ def _get_header_info(
         tokens_limit = int(tokens_limit_str) if tokens_limit_str is not None else None
         tokens_remaining_str = headers.get('x-ratelimit-remaining-tokens')
         tokens_remaining = int(tokens_remaining_str) if tokens_remaining_str is not None else None
-        tokens_reset_str = headers.get('x-ratelimit-reset-tokens')
+        tokens_reset_str = headers.get('x-ratelimit-reset-tokens', '5s')  # Default to 5 seconds
         tokens_reset_ts = now + _parse_header_duration(tokens_reset_str)
         tokens_info = (tokens_limit, tokens_remaining, tokens_reset_ts)
 
@@ -207,7 +210,7 @@ async def speech(input: str, *, model: str, voice: str, model_kwargs: Optional[d
 
     content = await _openai_client().audio.speech.create(input=input, model=model, voice=voice, **model_kwargs)
     ext = model_kwargs.get('response_format', 'mp3')
-    output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
+    output_filename = str(env.Env.get().create_tmp_path(f'.{ext}'))
     content.write_to_file(output_filename)
     return output_filename
 

pixeltable/functions/timestamp.py

@@ -237,12 +237,12 @@ def _(
     microsecond: sql.ColumnElement = _SQL_ZERO,
 ) -> sql.ColumnElement:
     return sql.func.make_timestamptz(
-        sql.cast(year, sql.Integer),
-        sql.cast(month, sql.Integer),
-        sql.cast(day, sql.Integer),
-        sql.cast(hour, sql.Integer),
-        sql.cast(minute, sql.Integer),
-        sql.cast(second + microsecond / 1000000.0, sql.Float),
+        year.cast(sql.Integer),
+        month.cast(sql.Integer),
+        day.cast(sql.Integer),
+        hour.cast(sql.Integer),
+        minute.cast(sql.Integer),
+        (second + microsecond / 1000000.0).cast(sql.Float),
     )
 
 

pixeltable/functions/video.py

@@ -2,9 +2,6 @@
 Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for `VideoType`.
 """
 
-import tempfile
-import uuid
-from pathlib import Path
 from typing import Any, Optional
 
 import av
@@ -59,8 +56,7 @@ class make_video(pxt.Aggregator):
         if frame is None:
             return
         if self.container is None:
-            (_, output_filename) = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
-            self.out_file = Path(output_filename)
+            self.out_file = env.Env.get().create_tmp_path('.mp4')
             self.container = av.open(str(self.out_file), mode='w')
             self.stream = self.container.add_stream('h264', rate=self.fps)
             self.stream.pix_fmt = 'yuv420p'
@@ -109,7 +105,7 @@ def extract_audio(
             return None
         audio_stream = container.streams.audio[stream_idx]
         # create this in our tmp directory, so it'll get cleaned up if it's being generated as part of a query
-        output_filename = str(env.Env.get().tmp_dir / f'{uuid.uuid4()}.{ext}')
+        output_filename = str(env.Env.get().create_tmp_path(f'.{ext}'))
 
         with av.open(output_filename, 'w', format=format) as output_container:
             output_stream = output_container.add_stream(codec or default_codec)