def test_cc_limit_to():
c = ColumnContainer(["a", "b", "c"])
c2 = c.limit_to(["c", "a"])
assert c2.columns == ["c", "a"]
assert c2.mapping() == [("a", "a"), ("b", "b"), ("c", "c")]
assert c.columns == ["a", "b", "c"]
assert c.mapping() == [("a", "a"), ("b", "b"), ("c", "c")]
def fix_column_to_row_type(
cc: ColumnContainer,
row_type: "org.apache.calcite.rel.type.RelDataType"
) -> ColumnContainer:
"""
Make sure that the given column container
has the column names specified by the row type.
We assume that the column order is already correct
and will just "blindly" rename the columns.
"""
field_names = [str(x) for x in row_type.getFieldNames()]
cc = cc.rename(columns=dict(zip(cc.columns, field_names)))
# TODO: We can also check for the types here and do any conversions if needed
return cc.limit_to(field_names)
