AB infer missing values on row

fill_missing_values.py

from typing import List
import pandas as pd

df = pd.DataFrame(
    [
        ["device_1", None, 1.0, 2.0, None, None, 3.0],
        ["device_2", 1.0, None, 2.0, None, None, None],
    ],
    columns=["device_id", "0", "1", "2", "3", "4", "5"],
)

print(df)


def fill_set_points(row: pd.Series, columns: List[str]) -> pd.Series:
    """Fill set points for a row.
    Assigns the first non-null value to the first null value.
    Warning: This function assumes that columns are ordered.

    :param: row (pd.Series): Row with null values
    :param: columns (List[str]): List of column names to fill null values

    Returns:
        pd.Series: Row with set points filled.
    """
    # find non-null indices and values among a set of columns
    non_null_indices = row[columns][row.notnull()].index
    last_non_null_values = row[columns][row.notnull()].values[0]

    # fill null values with the last non-null value for a set of columns
    for i in columns:
        if i not in non_null_indices:
            row[i] = last_non_null_values  # fill null
        else:
            last_non_null_values = row[i]  # update last known value

    return row


columns = [str(i) for i in range(6)]

df = df.apply(fill_set_points, axis=1, columns=columns)

print(df)