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

pixeltable/io/hf_datasets.py

@@ -10,6 +10,8 @@ import pixeltable as pxt
 import pixeltable.type_system as ts
 from pixeltable import exceptions as excs
 
+from .utils import normalize_import_parameters, normalize_schema_names
+
 if typing.TYPE_CHECKING:
     import datasets  # type: ignore[import-untyped]
 
@@ -28,29 +30,33 @@ _hf_to_pxt: dict[str, ts.ColumnType] = {
     'int64': ts.IntType(nullable=True),
     'bool': ts.BoolType(nullable=True),
     'float32': ts.FloatType(nullable=True),
+    'float64': ts.FloatType(nullable=True),
+    'large_string': ts.StringType(nullable=True),
     'string': ts.StringType(nullable=True),
     'timestamp[s]': ts.TimestampType(nullable=True),
     'timestamp[ms]': ts.TimestampType(nullable=True),  # HF dataset iterator converts timestamps to datetime.datetime
+    'timestamp[us]': ts.TimestampType(nullable=True),
 }
 
 
-def _to_pixeltable_type(feature_type: Any) -> Optional[ts.ColumnType]:
+def _to_pixeltable_type(feature_type: Any, nullable: bool) -> Optional[ts.ColumnType]:
     """Convert a huggingface feature type to a pixeltable ColumnType if one is defined."""
     import datasets
 
     if isinstance(feature_type, datasets.ClassLabel):
         # enum, example: ClassLabel(names=['neg', 'pos'], id=None)
-        return ts.StringType(nullable=True)
+        return ts.StringType(nullable=nullable)
     elif isinstance(feature_type, datasets.Value):
         # example: Value(dtype='int64', id=None)
-        return _hf_to_pxt.get(feature_type.dtype, None)
+        pt = _hf_to_pxt.get(feature_type.dtype, None)
+        return pt.copy(nullable=nullable) if pt is not None else None
     elif isinstance(feature_type, datasets.Sequence):
         # example: cohere wiki. Sequence(feature=Value(dtype='float32', id=None), length=-1, id=None)
-        dtype = _to_pixeltable_type(feature_type.feature)
+        dtype = _to_pixeltable_type(feature_type.feature, nullable)
         length = feature_type.length if feature_type.length != -1 else None
         return ts.ArrayType(shape=(length,), dtype=dtype)
     elif isinstance(feature_type, datasets.Image):
-        return ts.ImageType(nullable=True)
+        return ts.ImageType(nullable=nullable)
     else:
         return None
 
@@ -63,15 +69,17 @@ def _get_hf_schema(dataset: Union[datasets.Dataset, datasets.DatasetDict]) -> da
     return first_dataset.features
 
 
-def huggingface_schema_to_pixeltable_schema(
-    hf_dataset: Union[datasets.Dataset, datasets.DatasetDict],
+def huggingface_schema_to_pxt_schema(
+    hf_schema: datasets.Features, schema_overrides: dict[str, Any], primary_key: list[str]
 ) -> dict[str, Optional[ts.ColumnType]]:
     """Generate a pixeltable schema from a huggingface dataset schema.
     Columns without a known mapping are mapped to None
     """
-    hf_schema = _get_hf_schema(hf_dataset)
     pixeltable_schema = {
-        column_name: _to_pixeltable_type(feature_type) for column_name, feature_type in hf_schema.items()
+        column_name: _to_pixeltable_type(feature_type, column_name not in primary_key)
+        if column_name not in schema_overrides
+        else schema_overrides[column_name]
+        for column_name, feature_type in hf_schema.items()
     }
     return pixeltable_schema
 
@@ -82,6 +90,7 @@ def import_huggingface_dataset(
     *,
     column_name_for_split: Optional[str] = None,
     schema_overrides: Optional[dict[str, Any]] = None,
+    primary_key: Optional[Union[str, list[str]]] = None,
     **kwargs: Any,
 ) -> pxt.Table:
     """Create a new base table from a Huggingface dataset, or dataset dict with multiple splits.
@@ -97,6 +106,7 @@ def import_huggingface_dataset(
             name `name` will be given type `type`, instead of being inferred from the `Dataset` or `DatasetDict`. The keys in
             `schema_overrides` should be the column names of the `Dataset` or `DatasetDict` (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:
@@ -106,57 +116,47 @@ def import_huggingface_dataset(
 
     import pixeltable as pxt
 
-    if table_path in pxt.list_tables():
-        raise excs.Error(f'table {table_path} already exists')
-
     if not isinstance(dataset, (datasets.Dataset, datasets.DatasetDict)):
         raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
 
-    if isinstance(dataset, datasets.Dataset):
-        # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
-        raw_name = dataset.split._name
-        split_name = raw_name.split('[')[0] if raw_name is not None else None
-        dataset_dict = {split_name: dataset}
-    else:
-        dataset_dict = dataset
-
-    pixeltable_schema = huggingface_schema_to_pixeltable_schema(dataset)
-    if schema_overrides is not None:
-        pixeltable_schema.update(schema_overrides)
+    # Create the pixeltable schema from the huggingface schema
+    hf_schema_source = _get_hf_schema(dataset)
+    schema_overrides, primary_key = normalize_import_parameters(schema_overrides, primary_key)
+    hf_schema = huggingface_schema_to_pxt_schema(hf_schema_source, schema_overrides, primary_key)
 
+    # Add the split column to the schema if requested
     if column_name_for_split is not None:
-        if column_name_for_split in pixeltable_schema:
+        if column_name_for_split in hf_schema:
             raise excs.Error(
                 f'Column name `{column_name_for_split}` already exists in dataset schema; provide a different `column_name_for_split`'
             )
-        pixeltable_schema[column_name_for_split] = ts.StringType(nullable=True)
+        hf_schema[column_name_for_split] = ts.StringType(nullable=True)
 
-    for field, column_type in pixeltable_schema.items():
-        if column_type is None:
-            raise excs.Error(f'Could not infer pixeltable type for feature `{field}` in huggingface dataset')
+    schema, pxt_pk, _ = normalize_schema_names(hf_schema, primary_key, schema_overrides, True)
+
+    # Prepare to create table and insert data
+    if table_path in pxt.list_tables():
+        raise excs.Error(f'table {table_path} already exists')
 
     if isinstance(dataset, datasets.Dataset):
         # when loading an hf dataset partially, dataset.split._name is sometimes the form "train[0:1000]"
         raw_name = dataset.split._name
         split_name = raw_name.split('[')[0] if raw_name is not None else None
         dataset_dict = {split_name: dataset}
-    elif isinstance(dataset, datasets.DatasetDict):
-        dataset_dict = dataset
     else:
-        raise excs.Error(f'`type(dataset)` must be `datasets.Dataset` or `datasets.DatasetDict`. Got {type(dataset)=}')
+        dataset_dict = dataset
 
     # extract all class labels from the dataset to translate category ints to strings
-    hf_schema = _get_hf_schema(dataset)
     categorical_features = {
         feature_name: feature_type.names
-        for (feature_name, feature_type) in hf_schema.items()
+        for (feature_name, feature_type) in hf_schema_source.items()
         if isinstance(feature_type, datasets.ClassLabel)
     }
 
     try:
         # random tmp name
         tmp_name = f'{table_path}_tmp_{random.randint(0, 100000000)}'
-        tab = pxt.create_table(tmp_name, pixeltable_schema, **kwargs)
+        tab = pxt.create_table(tmp_name, schema, primary_key=pxt_pk, **kwargs)
 
         def _translate_row(row: dict[str, Any], split_name: str) -> dict[str, Any]:
             output_row = row.copy()

pixeltable/io/pandas.py

@@ -2,17 +2,21 @@ from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
+from pandas._typing import DtypeObj  # For pandas dtype type hints
+from pandas.api.types import is_datetime64_any_dtype, is_extension_array_dtype
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
-import pixeltable.type_system as ts
+from pixeltable import Table
+
+from .utils import find_or_create_table, normalize_import_parameters, normalize_schema_names
 
 
 def import_pandas(
     tbl_name: str,
     df: pd.DataFrame,
     *,
-    schema_overrides: Optional[dict[str, pxt.ColumnType]] = None,
+    schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
@@ -39,16 +43,16 @@ def import_pandas(
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
     """
-    if schema_overrides is None:
-        schema_overrides = {}
-    if primary_key is None:
-        primary_key = []
-    elif isinstance(primary_key, str):
-        primary_key = [primary_key]
-
-    schema, pxt_pk = __df_to_pxt_schema(df, schema_overrides, primary_key)
-    tbl_rows = (dict(__df_row_to_pxt_row(row, schema)) for row in df.itertuples())
-    table = pxt.create_table(
+    schema_overrides, primary_key = normalize_import_parameters(schema_overrides, primary_key)
+    pd_schema = df_infer_schema(df, schema_overrides, primary_key)
+    schema, pxt_pk, col_mapping = normalize_schema_names(pd_schema, primary_key, schema_overrides, False)
+
+    __check_primary_key_values(df, primary_key)
+
+    # Convert all rows to insertable format
+    tbl_rows = [__df_row_to_pxt_row(row, pd_schema, col_mapping) for row in df.itertuples()]
+
+    table = find_or_create_table(
         tbl_name, schema, primary_key=pxt_pk, num_retained_versions=num_retained_versions, comment=comment
     )
     table.insert(tbl_rows)
@@ -58,7 +62,7 @@ def import_pandas(
 def import_csv(
     tbl_name: str,
     filepath_or_buffer,
-    schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
+    schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
@@ -88,7 +92,7 @@ def import_excel(
     tbl_name: str,
     io,
     *args,
-    schema_overrides: Optional[dict[str, ts.ColumnType]] = None,
+    schema_overrides: Optional[dict[str, Any]] = None,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
@@ -114,82 +118,73 @@ def import_excel(
     )
 
 
-def __df_to_pxt_schema(
+def __check_primary_key_values(df: pd.DataFrame, primary_key: list[str]) -> None:
+    for pd_name in primary_key:
+        # This can be faster for large DataFrames
+        has_nulls = df[pd_name].count() < len(df)
+        if has_nulls:
+            raise excs.Error(f'Primary key column `{pd_name}` cannot contain null values.')
+
+
+def df_infer_schema(
     df: pd.DataFrame, schema_overrides: dict[str, pxt.ColumnType], primary_key: list[str]
-) -> tuple[dict[str, pxt.ColumnType], list[str]]:
+) -> dict[str, pxt.ColumnType]:
     """
     Infers a Pixeltable schema from a Pandas DataFrame.
 
     Returns:
         A tuple containing a Pixeltable schema and a list of primary key column names.
     """
-    for pd_name in schema_overrides:
-        if pd_name not in df.columns:
-            raise excs.Error(
-                f'Column `{pd_name}` specified in `schema_overrides` does not exist in the given `DataFrame`.'
-            )
-    for pd_name in primary_key:
-        if pd_name not in df.columns:
-            raise excs.Error(f'Primary key column `{pd_name}` does not exist in the given `DataFrame`.')
-
-    schema: dict[str, pxt.ColumnType] = {}
-    col_mapping: dict[str, str] = {}  # Maps Pandas column names to Pixeltable column names
-
+    pd_schema: dict[str, pxt.ColumnType] = {}
     for pd_name, pd_dtype in zip(df.columns, df.dtypes):
         if pd_name in schema_overrides:
             pxt_type = schema_overrides[pd_name]
         else:
-            # This complicated-looking condition is necessary because we cannot safely call `pd.isna()` on
-            # general objects, so we need to check for nulls in the specific cases where we might expect them.
-            # isinstance(val, float) will check for NaN values in float columns *as well as* floats appearing
-            # in object columns (where Pandas uses NaN as a general null).
-            # np.issubdtype(pd_dtype, np.datetime64) checks for NaT values specifically in datetime columns.
-            has_na = any(
-                (isinstance(val, float) or np.issubdtype(pd_dtype, np.datetime64)) and pd.isna(val)
-                for val in df[pd_name]
-            )
-            if has_na and pd_name in primary_key:
-                raise excs.Error(f'Primary key column `{pd_name}` cannot contain null values.')
-            pxt_type = __np_dtype_to_pxt_type(pd_dtype, df[pd_name], pd_name not in primary_key)
-        pxt_name = __normalize_pxt_col_name(pd_name)
-        # Ensure that column names are unique by appending a distinguishing suffix
-        # to any collisions
-        if pxt_name in schema:
-            n = 2
-            while f'{pxt_name}_{n}' in schema:
-                n += 1
-            pxt_name = f'{pxt_name}_{n}'
-        schema[pxt_name] = pxt_type
-        col_mapping[pd_name] = pxt_name
-
-    pxt_pk = [col_mapping[pk] for pk in primary_key]
-    return schema, pxt_pk
-
-
-def __normalize_pxt_col_name(pd_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
+            pxt_type = __pd_coltype_to_pxt_type(pd_dtype, df[pd_name], pd_name not in primary_key)
+        pd_schema[pd_name] = pxt_type
+
+    return pd_schema
+
+
+"""
+# Check if a datetime64[ns, UTC] dtype
+def is_datetime_tz_utc(x: Any) -> bool:
+    if isinstance(x, pd.Timestamp) and x.tzinfo is not None and str(x.tzinfo) == 'UTC':
+        return True
+    return pd.api.types.is_datetime64tz_dtype(x) and str(x).endswith('UTC]')
+"""
+
+
+def __pd_dtype_to_pxt_type(pd_dtype: DtypeObj, nullable: bool) -> Optional[pxt.ColumnType]:
     """
-    id = ''.join(ch if ch.isascii() and ch.isalnum() else '_' for ch in pd_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
+    Determines a pixeltable ColumnType from a pandas dtype
 
+    Args:
+        pd_dtype: A pandas dtype object
 
-def __np_dtype_to_pxt_type(np_dtype: np.dtype, data_col: pd.Series, nullable: bool) -> pxt.ColumnType:
+    Returns:
+        pxt.ColumnType: A pixeltable ColumnType
+    """
+    # Pandas extension arrays / types (Int64, boolean, string[pyarrow], etc.) are not directly compatible with NumPy dtypes
+    # The timezone-aware datetime64[ns, tz=] dtype is a pandas extension dtype
+    if is_datetime64_any_dtype(pd_dtype):
+        return pxt.TimestampType(nullable=nullable)
+    if is_extension_array_dtype(pd_dtype):
+        return None
+    # Most other pandas dtypes are directly NumPy compatible
+    assert isinstance(pd_dtype, np.dtype)
+    return pxt.ArrayType.from_np_dtype(pd_dtype, nullable)
+
+
+def __pd_coltype_to_pxt_type(pd_dtype: DtypeObj, data_col: pd.Series, nullable: bool) -> pxt.ColumnType:
     """
-    Infers a Pixeltable type based on a Numpy dtype.
+    Infers a Pixeltable type based on a pandas dtype.
     """
-    pxttype = ts.ArrayType.from_np_dtype(np_dtype, nullable)
+    pxttype = __pd_dtype_to_pxt_type(pd_dtype, nullable)
     if pxttype is not None:
         return pxttype
 
-    if np_dtype == np.object_:
+    if pd_dtype == np.object_:
         # The `object_` dtype can mean all sorts of things; see if we can infer the Pixeltable type
         # based on the actual data in `data_col`.
         # First drop any null values (they don't contribute to type inference).
@@ -206,11 +201,14 @@ def __np_dtype_to_pxt_type(np_dtype: np.dtype, data_col: pd.Series, nullable: bo
         else:
             return inferred_type.copy(nullable=nullable)
 
-    raise excs.Error(f'Could not infer Pixeltable type of column: {data_col.name} (dtype: {np_dtype})')
+    raise excs.Error(f'Could not infer Pixeltable type of column: {data_col.name} (dtype: {pd_dtype})')
 
 
-def __df_row_to_pxt_row(row: tuple[Any, ...], schema: dict[str, pxt.ColumnType]) -> dict[str, Any]:
-    rows = {}
+def __df_row_to_pxt_row(
+    row: tuple[Any, ...], schema: dict[str, pxt.ColumnType], col_mapping: Optional[dict[str, str]]
+) -> dict[str, Any]:
+    """Convert a row to insertable format"""
+    pxt_row: dict[str, Any] = {}
     for val, (col_name, pxt_type) in zip(row[1:], schema.items()):
         if pxt_type.is_float_type():
             val = float(val)
@@ -232,5 +230,6 @@ def __df_row_to_pxt_row(row: tuple[Any, ...], schema: dict[str, pxt.ColumnType])
                 val = None
             else:
                 val = pd.Timestamp(val).to_pydatetime()
-        rows[col_name] = val
-    return rows
+        pxt_name = col_name if col_mapping is None else col_mapping[col_name]
+        pxt_row[pxt_name] = val
+    return pxt_row

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