What causes "Expected 1d array, got array with shape {col.shape}" in mlflow/mlflow?

Reshape each column to be strictly 1D: flatten or keep per-cell values as ragged lists so col.values.shape has ndim 1.

mlflow/mlflow · error · MlflowException

Expected 1d array, got array with shape {col.shape}

Error message

Expected 1d array, got array with shape {col.shape}

What it means

MLflow's schema inference (_infer_pandas_column) requires each pandas DataFrame column to be a 1-dimensional pd.Series. This error is thrown when the underlying values of a column have more than one dimension (e.g. a column of arrays/vectors), because MLflow's Schema types (DataType, TensorSpec) cannot represent a 2D object column. It protects against silently inferring a wrong model signature.

Source

Thrown at mlflow/types/utils.py:512

def _is_none_or_nan(x):
    if isinstance(x, float):
        return np.isnan(x)
    # NB: We can't use pd.isna() because the input can be a series.
    return x is None or x is pd.NA or x is pd.NaT


def _infer_required(col) -> bool:
    if isinstance(col, (list, pd.Series)):
        return not any(_is_none_or_nan(x) for x in col)
    return not _is_none_or_nan(col)


def _infer_pandas_column(col: pd.Series) -> DataType:
    if not isinstance(col, pd.Series):
        raise TypeError(f"Expected pandas.Series, got '{type(col)}'.")
    if len(col.values.shape) > 1:
        raise MlflowException(f"Expected 1d array, got array with shape {col.shape}")

    if col.dtype.kind == "O":
        col = col.infer_objects()
    if col.dtype.kind == "O":
        try:
            # We convert pandas Series into list and infer the schema.
            # The real schema for internal field should be the Array's dtype
            arr_type = _infer_colspec_type(col.to_list())
            return arr_type.dtype
        except Exception as e:
            # For backwards compatibility, we fall back to string
            # if the provided array is of string type
            if pd.api.types.is_string_dtype(col):
                return DataType.string
            raise MlflowException(f"Failed to infer schema for pandas.Series {col}. Error: {e}")
    else:
        # NB: The following works for numpy types as well as pandas extension types.
        return _infer_numpy_dtype(col.dtype)

View on GitHub (pinned to 6a27f2decc)

Solutions

Reshape each column to be strictly 1D: flatten or keep per-cell values as ragged lists so col.values.shape has ndim 1.
For tensor/multi-dim data, use mlflow.models.infer_signature with a numpy ndarray input instead of a pandas DataFrame, so a TensorSpec is inferred instead.
Convert per-cell arrays into separate scalar columns, or into lists (object dtype with 1D shape).
Cast the offending column explicitly: df['x'] = df['x'].apply(lambda v: list(v)) and re-run inference.

Example fix

// before
df = pd.DataFrame({'x': np.random.rand(10, 3, 3)})  # column values shape (10,3,3)
sig = infer_signature(df, model)
// after
df = pd.DataFrame({'x': [np.random.rand(3, 3).tolist() for _ in range(10)]})  # object col of lists
# or pass ndarray input to get a TensorSpec:
sig = infer_signature(np.random.rand(10, 3, 3), model)

Defensive patterns

Strategy: validation

Validate before calling

import pandas as pd
def check_1d_columns(df: pd.DataFrame):
    bad = [name for name in df.columns if getattr(df[name].values, 'ndim', 1) > 1]
    if bad:
        raise ValueError(f'Columns with >1D values: {bad}')

Type guard

def is_1d_series(col) -> bool:
    import pandas as pd
    return isinstance(col, pd.Series) and col.values.ndim == 1

Prevention

Keep DataFrame columns scalar or ragged-list (object) typed; never store rectangular arrays as column values.
Use ndarray inputs (not DataFrames) for tensor data so TensorSpec is inferred.
Run signature inference early in training, not at deployment time.
Add a column-ndim assertion in data-preparation tests.

When it happens

Trigger: Calling mlflow.models.infer_signature(model, pandas_df) (or _infer_schema directly) where a DataFrame column holds multi-dimensional data, e.g. df['x'] = np.random.rand(10, 3, 3) (each cell a 2D array) so col.values.shape is (n, 3, 3). Also happens when a column was built from lists of unequal arrays that pandas stored as object arrays with ndim > 1.

Common situations: Passing DataFrames with image arrays, embedding matrices, or multi-dim feature tensors as columns; converting numpy arrays with np.stack incorrectly; pandas columns created from lists of numpy arrays that end up rectangular (shape (n, m)).

Related errors

AI-assisted analysis of mlflow/mlflow@6a27f2decc (2026-08-29). Data as JSON: /api/errors/b4a11330fcf8d4bf. Report an issue: GitHub.