pixeltable 0.3.4__py3-none-any.whl → 0.3.6__py3-none-any.whl

pixeltable/io/parquet.py

@@ -15,10 +15,11 @@ import PIL.Image
 
 import pixeltable as pxt
 import pixeltable.exceptions as exc
-import pixeltable.type_system as ts
 from pixeltable.env import Env
 from pixeltable.utils.transactional_directory import transactional_directory
 
+from .utils import normalize_import_parameters, normalize_schema_names
+
 if typing.TYPE_CHECKING:
     import pyarrow as pa
 
@@ -148,19 +149,13 @@ def export_parquet(
         _write_batch(current_value_batch, arrow_schema, temp_path / f'part-{batch_num:05d}.parquet')
 
 
-def parquet_schema_to_pixeltable_schema(parquet_path: str) -> dict[str, Optional[ts.ColumnType]]:
-    """Generate a default pixeltable schema for the given parquet file. Returns None for unknown types."""
-    from pyarrow import parquet
-
-    from pixeltable.utils.arrow import to_pixeltable_schema
-
-    input_path = Path(parquet_path).expanduser()
-    parquet_dataset = parquet.ParquetDataset(str(input_path))
-    return to_pixeltable_schema(parquet_dataset.schema)
-
-
 def import_parquet(
-    table: str, *, parquet_path: str, schema_overrides: Optional[dict[str, ts.ColumnType]] = None, **kwargs: Any
+    table: str,
+    *,
+    parquet_path: str,
+    schema_overrides: Optional[dict[str, Any]] = None,
+    primary_key: Optional[Union[str, list[str]]] = None,
+    **kwargs: Any,
 ) -> pxt.Table:
     """Creates a new base table from a Parquet file or set of files. Requires pyarrow to be installed.
 
@@ -171,6 +166,7 @@ def import_parquet(
             name `name` will be given type `type`, instead of being inferred from the Parquet dataset. The keys in
             `schema_overrides` should be the column names of the Parquet dataset (whether or not they are valid
             Pixeltable identifiers).
+        primary_key: The primary key of the table (see [`create_table()`][pixeltable.create_table]).
         kwargs: Additional arguments to pass to `create_table`.
 
     Returns:
@@ -178,33 +174,29 @@ def import_parquet(
     """
     from pyarrow import parquet
 
-    import pixeltable as pxt
-    from pixeltable.utils.arrow import iter_tuples
+    from pixeltable.utils.arrow import ar_infer_schema, iter_tuples2
 
     input_path = Path(parquet_path).expanduser()
     parquet_dataset = parquet.ParquetDataset(str(input_path))
 
-    schema = parquet_schema_to_pixeltable_schema(parquet_path)
-    if schema_overrides is None:
-        schema_overrides = {}
-
-    schema.update(schema_overrides)
-    for k, v in schema.items():
-        if v is None:
-            raise exc.Error(f'Could not infer pixeltable type for column {k} from parquet file')
+    schema_overrides, primary_key = normalize_import_parameters(schema_overrides, primary_key)
+    ar_schema = ar_infer_schema(parquet_dataset.schema, schema_overrides, primary_key)
+    schema, pxt_pk, col_mapping = normalize_schema_names(ar_schema, primary_key, schema_overrides, False)
 
     if table in pxt.list_tables():
         raise exc.Error(f'Table {table} already exists')
 
+    tmp_name = f'{table}_tmp_{random.randint(0, 100000000)}'
+    total_rows = 0
     try:
-        tmp_name = f'{table}_tmp_{random.randint(0, 100000000)}'
-        tab = pxt.create_table(tmp_name, schema, **kwargs)
+        tab = pxt.create_table(tmp_name, schema, primary_key=pxt_pk, **kwargs)
         for fragment in parquet_dataset.fragments:  # type: ignore[attr-defined]
             for batch in fragment.to_batches():
-                dict_batch = list(iter_tuples(batch))
+                dict_batch = list(iter_tuples2(batch, col_mapping, schema))
+                total_rows += len(dict_batch)
                 tab.insert(dict_batch)
     except Exception as e:
-        _logger.error(f'Error while inserting Parquet file into table: {e}')
+        _logger.error(f'Error after inserting {total_rows} rows from Parquet file into table: {e}')
         raise e
 
     pxt.move(tmp_name, table)

pixeltable/io/utils.py

@@ -0,0 +1,115 @@
+from keyword import iskeyword as is_python_keyword
+from typing import Any, Optional, Union
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable import Table
+from pixeltable.catalog.globals import is_system_column_name
+
+
+def normalize_pxt_col_name(name: str) -> str:
+    """
+    Normalizes an arbitrary DataFrame column name into a valid Pixeltable identifier by:
+    - replacing any non-ascii or non-alphanumeric characters with an underscore _
+    - prefixing the result with the letter 'c' if it starts with an underscore or a number
+    """
+    id = ''.join(ch if ch.isascii() and ch.isalnum() else '_' for ch in name)
+    if id[0].isnumeric():
+        id = f'c_{id}'
+    elif id[0] == '_':
+        id = f'c{id}'
+    assert pxt.catalog.is_valid_identifier(id), id
+    return id
+
+
+def normalize_import_parameters(
+    schema_overrides: Optional[dict[str, Any]] = None, primary_key: Optional[Union[str, list[str]]] = None
+) -> tuple[dict[str, Any], list[str]]:
+    if schema_overrides is None:
+        schema_overrides = {}
+    if primary_key is None:
+        primary_key = []
+    elif isinstance(primary_key, str):
+        primary_key = [primary_key]
+    return schema_overrides, primary_key
+
+
+def _is_usable_as_column_name(name: str, destination_schema: dict[str, Any]) -> bool:
+    return not (is_system_column_name(name) or is_python_keyword(name) or name in destination_schema)
+
+
+def normalize_schema_names(
+    in_schema: dict[str, Any],
+    primary_key: list[str],
+    schema_overrides: dict[str, Any],
+    require_valid_pxt_column_names: bool = False,
+) -> tuple[dict[str, Any], list[str], Optional[dict[str, str]]]:
+    """
+    Convert all names in the input schema from source names to valid Pixeltable identifiers
+    - Ensure that all names are unique.
+    - Report an error if any types are missing
+    - If "require_valid_pxt_column_names", report an error if any column names are not valid Pixeltable column names
+    - Report an error if any primary key columns are missing
+    Returns
+    - A new schema with normalized column names
+    - The primary key columns, mapped to the normalized names
+    - A mapping from the original names to the normalized names.
+    """
+
+    # Report any untyped columns as an error
+    untyped_cols = [in_name for in_name, column_type in in_schema.items() if column_type is None]
+    if len(untyped_cols) > 0:
+        raise excs.Error(f'Could not infer pixeltable type for column(s): {", ".join(untyped_cols)}')
+
+    # Report any columns in `schema_overrides` that are not in the source
+    extraneous_overrides = schema_overrides.keys() - in_schema.keys()
+    if len(extraneous_overrides) > 0:
+        raise excs.Error(
+            f'Some column(s) specified in `schema_overrides` are not present in the source: {", ".join(extraneous_overrides)}'
+        )
+
+    schema: dict[str, Any] = {}
+    col_mapping: dict[str, str] = {}  # Maps column names to Pixeltable column names if needed
+    for in_name, pxt_type in in_schema.items():
+        pxt_name = normalize_pxt_col_name(in_name)
+        # Ensure that column names are unique by appending a distinguishing suffix
+        # to any collisions
+        pxt_fname = pxt_name
+        n = 1
+        while not _is_usable_as_column_name(pxt_fname, schema):
+            pxt_fname = f'{pxt_name}_{n}'
+            n += 1
+        schema[pxt_fname] = pxt_type
+        col_mapping[in_name] = pxt_fname
+
+    # Determine if the col_mapping is the identity mapping
+    non_identity_keys = [k for k, v in col_mapping.items() if k != v]
+    if len(non_identity_keys) > 0:
+        if require_valid_pxt_column_names:
+            raise excs.Error(
+                f'Column names must be valid pixeltable identifiers. Invalid names: {", ".join(non_identity_keys)}'
+            )
+    else:
+        col_mapping = None
+
+    # Report any primary key columns that are not in the source as an error
+    missing_pk = [pk for pk in primary_key if pk not in in_schema]
+    if len(missing_pk) > 0:
+        raise excs.Error(f'Primary key column(s) are not found in the source: {", ".join(missing_pk)}')
+
+    pxt_pk = [col_mapping[pk] for pk in primary_key] if col_mapping is not None else primary_key
+
+    return schema, pxt_pk, col_mapping
+
+
+def find_or_create_table(
+    tbl_path: str,
+    schema: dict[str, Any],
+    *,
+    primary_key: Optional[Union[str, list[str]]],
+    num_retained_versions: int,
+    comment: str,
+) -> Table:
+    return pxt.create_table(
+        tbl_path, schema, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment
+    )

pixeltable/iterators/audio.py

@@ -5,7 +5,7 @@ from fractions import Fraction
 from pathlib import Path
 from typing import Any, Optional
 
-import av  # type: ignore[import-untyped]
+import av
 
 import pixeltable.env as env
 import pixeltable.exceptions as excs
@@ -146,6 +146,7 @@ class AudioSplitter(ComponentIterator):
         input_stream = self.container.streams.audio[0]
         codec_name = AudioSplitter.__codec_map.get(input_stream.codec_context.name, input_stream.codec_context.name)
         output_stream = output_container.add_stream(codec_name, rate=input_stream.codec_context.sample_rate)
+        assert isinstance(output_stream, av.audio.stream.AudioStream)
         frame_count = 0
         # Since frames don't align with chunk boundaries, we may have read an extra frame in previous iteration
         # Seek to the nearest frame in stream at current chunk start time

pixeltable/iterators/video.py

@@ -4,7 +4,7 @@ from fractions import Fraction
 from pathlib import Path
 from typing import Any, Optional, Sequence
 
-import av  # type: ignore[import-untyped]
+import av
 import pandas as pd
 import PIL.Image
 

pixeltable/metadata/__init__.py

@@ -10,7 +10,7 @@ import sqlalchemy.orm as orm
 from .schema import SystemInfo, SystemInfoMd
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 27
+VERSION = 30
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:
@@ -31,6 +31,7 @@ converter_cbs: dict[int, Callable[[sql.engine.Engine], None]] = {}
 def register_converter(version: int) -> Callable[[Callable[[sql.engine.Engine], None]], None]:
     def decorator(fn: Callable[[sql.engine.Engine], None]) -> None:
         global converter_cbs
+        assert version not in converter_cbs
         converter_cbs[version] = fn
 
     return decorator

pixeltable/metadata/converters/convert_15.py

@@ -17,7 +17,7 @@ _logger = logging.getLogger('pixeltable')
 def _(engine: sql.engine.Engine) -> None:
     with engine.begin() as conn:
         for row in conn.execute(sql.select(Function)):
-            id, dir_id, md, binary_obj = row
+            id, _, md, binary_obj = row
             md['md'] = __update_md(md['md'], binary_obj)
             _logger.info(f'Updating function: {id}')
             conn.execute(sql.update(Function).where(Function.id == id).values(md=md))
@@ -27,14 +27,24 @@ def __update_md(orig_d: dict, binary_obj: bytes) -> Any:
     # construct dict produced by CallableFunction.to_store()
     py_fn = cloudpickle.loads(binary_obj)
     py_params = inspect.signature(py_fn).parameters
-    return_type = ts.ColumnType.from_dict(orig_d['return_type'])
-    params: list[func.Parameter] = []
+    return_type = orig_d['return_type']
+    params: list[dict] = []
     for name, col_type_dict, kind_int, is_batched in orig_d['parameters']:
-        col_type = ts.ColumnType.from_dict(col_type_dict) if col_type_dict is not None else None
         default = py_params[name].default
-        kind = inspect._ParameterKind(kind_int)  # is there a way to avoid referencing a private type?
-        params.append(func.Parameter(name=name, col_type=col_type, kind=kind, default=default, is_batched=is_batched))
+        kind = inspect._ParameterKind(kind_int)
+        params.append(
+            {
+                'name': name,
+                'col_type': col_type_dict,
+                'kind': str(kind),
+                'is_batched': is_batched,
+                'has_default': default is not inspect.Parameter.empty,
+                'default': None if default is inspect.Parameter.empty else default,
+            }
+        )
     is_batched = 'batch_size' in orig_d
-    sig = func.Signature(return_type, params, is_batched=is_batched)
-    d = {'signature': sig.as_dict(), 'batch_size': orig_d['batch_size'] if is_batched else None}
+    d = {
+        'signature': {'return_type': return_type, 'parameters': params, 'is_batched': is_batched},
+        'batch_size': orig_d['batch_size'] if is_batched else None,
+    }
     return d

pixeltable/metadata/converters/convert_27.py

@@ -0,0 +1,31 @@
+import logging
+from typing import Any, Optional
+from uuid import UUID
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+from pixeltable.metadata.schema import Table
+
+_logger = logging.getLogger('pixeltable')
+
+
+@register_converter(version=27)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, table_md_updater=__update_table_md)
+
+
+def __update_table_md(table_md: dict, table_id: UUID) -> None:
+    """Update the view metadata to add the include_base_columns boolean if it is missing
+
+    Args:
+        table_md (dict): copy of the original table metadata. this gets updated in place.
+        table_id (UUID): the table id
+
+    """
+    if table_md['view_md'] is None:
+        return
+    if 'include_base_columns' not in table_md['view_md']:
+        table_md['view_md']['include_base_columns'] = True
+        _logger.info(f'Updating view metadata for table: {table_id}')

pixeltable/metadata/converters/convert_28.py

@@ -0,0 +1,15 @@
+import logging
+
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.schema import Dir, Table, TableSchemaVersion, TableVersion
+
+
+@register_converter(version=28)
+def _(engine: sql.engine.Engine) -> None:
+    with engine.begin() as conn:
+        conn.execute(sql.update(Dir).values(md=Dir.md.concat({'user': None, 'additional_md': {}})))
+        conn.execute(sql.update(Table).values(md=Table.md.concat({'user': None, 'additional_md': {}})))
+        conn.execute(sql.update(TableVersion).values(md=TableVersion.md.concat({'additional_md': {}})))
+        conn.execute(sql.update(TableSchemaVersion).values(md=TableSchemaVersion.md.concat({'additional_md': {}})))

pixeltable/metadata/converters/convert_29.py

@@ -0,0 +1,111 @@
+from typing import Any, Optional
+
+import sqlalchemy as sql
+
+from pixeltable import exprs
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_md
+
+
+@register_converter(version=29)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_md(engine, substitution_fn=__substitute_md)
+
+
+def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
+    # Defaults are now stored as literals in signatures
+    if k == 'parameters':
+        for param in v:
+            assert isinstance(param, dict)
+            has_default = param.get('has_default') or (param.get('default') is not None)
+            if 'has_default' in param:
+                del param['has_default']
+            literal = exprs.Expr.from_object(param['default']) if has_default else None
+            assert literal is None or isinstance(literal, exprs.Literal)
+            param['default'] = None if literal is None else literal.as_dict()
+        return k, v
+
+    # Method of organizing argument expressions has changed
+    if isinstance(v, dict) and v.get('_classname') == 'FunctionCall':
+        args = v['args']
+        kwargs = v['kwargs']
+        components = v['components']
+        group_by_start_idx = v['group_by_start_idx']
+        group_by_stop_idx = v['group_by_stop_idx']
+        order_by_start_idx = v['order_by_start_idx']
+
+        new_args = []
+        for arg in args:
+            if arg[0] is not None:
+                assert isinstance(arg[0], int)
+                new_args.append(components[arg[0]])
+            else:
+                literal = exprs.Expr.from_object(arg[1])
+                new_args.append(literal.as_dict())
+
+        new_kwargs = {}
+        for name, kwarg in kwargs.items():
+            if kwarg[0] is not None:
+                assert isinstance(kwarg[0], int)
+                new_kwargs[name] = components[kwarg[0]]
+            else:
+                literal = exprs.Expr.from_object(kwarg[1])
+                new_kwargs[name] = literal.as_dict()
+
+        # We need to expand ("unroll") any var-args or var-kwargs.
+
+        new_args_len = len(new_args)
+        rolled_args: Optional[dict] = None
+        rolled_kwargs: Optional[dict] = None
+
+        if 'signature' in v['fn']:
+            # If it's a pickled function, there's no signature, so we're out of luck; varargs in a pickled function
+            # is an edge case that won't migrate properly.
+            parameters: list[dict] = v['fn']['signature']['parameters']
+            for i, param in enumerate(parameters):
+                if param['kind'] == 'VAR_POSITIONAL':
+                    if new_args_len > i:
+                        # For peculiar historical reasons, variable kwargs might show up in args. Thus variable
+                        # positional args is not necessarily the last element of args; it might be the second-to-last.
+                        assert new_args_len <= i + 2, new_args
+                        rolled_args = new_args[i]
+                        new_args = new_args[:i] + new_args[i + 1 :]
+                if param['kind'] == 'VAR_KEYWORD':
+                    # As noted above, variable kwargs might show up either in args or in kwargs. If it's in args, it
+                    # is necessarily the last element.
+                    if new_args_len > i:
+                        assert new_args_len <= i + 1, new_args
+                        rolled_kwargs = new_args.pop()
+                    if param['name'] in kwargs:
+                        assert rolled_kwargs is None
+                        rolled_kwargs = kwargs.pop(param['name'])
+
+        if rolled_args is not None:
+            assert rolled_args['_classname'] in ('InlineArray', 'InlineList')
+            new_args.extend(rolled_args['components'])
+        if rolled_kwargs is not None:
+            assert rolled_kwargs['_classname'] == 'InlineDict'
+            new_kwargs.update(zip(rolled_kwargs['keys'], rolled_kwargs['components']))
+
+        group_by_exprs = [components[i] for i in range(group_by_start_idx, group_by_stop_idx)]
+        order_by_exprs = [components[i] for i in range(order_by_start_idx, len(components))]
+
+        new_components = [*new_args, *new_kwargs.values(), *group_by_exprs, *order_by_exprs]
+
+        newv = {
+            'fn': v['fn'],
+            'arg_idxs': list(range(len(new_args))),
+            'kwarg_idxs': {name: i + len(new_args) for i, name in enumerate(new_kwargs.keys())},
+            'group_by_start_idx': len(new_args) + len(new_kwargs),
+            'group_by_stop_idx': len(new_args) + len(new_kwargs) + len(group_by_exprs),
+            'order_by_start_idx': len(new_args) + len(new_kwargs) + len(group_by_exprs),
+            'is_method_call': False,
+            '_classname': 'FunctionCall',
+            'components': new_components,
+        }
+        if 'return_type' in v:
+            newv['return_type'] = v['return_type']
+
+        return k, newv
+
+    return None

pixeltable/metadata/converters/util.py

@@ -5,7 +5,7 @@ from uuid import UUID
 
 import sqlalchemy as sql
 
-from pixeltable.metadata.schema import Table, TableSchemaVersion
+from pixeltable.metadata.schema import Function, Table, TableSchemaVersion
 
 __logger = logging.getLogger('pixeltable')
 
@@ -50,6 +50,17 @@ def convert_table_md(
                 __logger.info(f'Updating schema for table: {id}')
                 conn.execute(sql.update(Table).where(Table.id == id).values(md=updated_table_md))
 
+        for row in conn.execute(sql.select(Function)):
+            id = row[0]
+            function_md = row[2]
+            assert isinstance(function_md, dict)
+            updated_function_md = copy.deepcopy(function_md)
+            if substitution_fn is not None:
+                updated_function_md = __substitute_md_rec(updated_function_md, substitution_fn)
+            if updated_function_md != function_md:
+                __logger.info(f'Updating function: {id}')
+                conn.execute(sql.update(Function).where(Function.id == id).values(md=updated_function_md))
+
 
 def __update_column_md(table_md: dict, column_md_updater: Callable[[dict], None]) -> None:
     columns_md = table_md['column_md']

pixeltable/metadata/notes.py

@@ -2,6 +2,9 @@
 # rather than as a comment, so that the existence of a description can be enforced by
 # the unit tests when new versions are added.
 VERSION_NOTES = {
+    30: 'Store default values and constant arguments as literals',
+    29: 'Add user and additional_md fields to metadata structs',
+    28: 'Enable view creation from DataFrame with select clause',
     27: 'Enable pxt.query parameterization of limit clauses',
     26: 'Rename clip_text and clip_image to clip',
     25: 'Functions with multiple signatures',

pixeltable/metadata/schema.py

@@ -74,6 +74,8 @@ class SystemInfo(Base):
 @dataclasses.dataclass
 class DirMd:
     name: str
+    user: Optional[str]
+    additional_md: dict[str, Any]
 
 
 class Dir(Base):
@@ -132,6 +134,7 @@ class IndexMd:
 @dataclasses.dataclass
 class ViewMd:
     is_snapshot: bool
+    include_base_columns: bool
 
     # (table id, version); for mutable views, all versions are None
     base_versions: list[tuple[str, Optional[int]]]
@@ -150,6 +153,8 @@ class ViewMd:
 class TableMd:
     name: str
 
+    user: Optional[str]
+
     # monotonically increasing w/in Table for both data and schema changes, starting at 0
     current_version: int
     # each version has a corresponding schema version (current_version >= current_schema_version)
@@ -169,6 +174,7 @@ class TableMd:
     column_md: dict[int, ColumnMd]  # col_id -> ColumnMd
     index_md: dict[int, IndexMd]  # index_id -> IndexMd
     view_md: Optional[ViewMd]
+    additional_md: dict[str, Any]
 
 
 class Table(Base):
@@ -194,6 +200,7 @@ class TableVersionMd:
     created_at: float  # time.time()
     version: int
     schema_version: int
+    additional_md: dict[str, Any]
 
 
 class TableVersion(Base):
@@ -232,6 +239,7 @@ class TableSchemaVersionMd:
     # default validation strategy for any media column of this table
     # stores column.MediaValiation.name.lower()
     media_validation: str
+    additional_md: dict[str, Any]
 
 
 # versioning: each table schema change results in a new record

pixeltable/share/__init__.py

@@ -0,0 +1 @@
+from .publish import publish_snapshot