Skip to content

Commit

Permalink
fix(memtables): disallow duplicate column names when constructing mem…
Browse files Browse the repository at this point in the history
…tables
  • Loading branch information
cpcloud authored and gforsyth committed Dec 7, 2023
1 parent d8d622c commit 4937b48
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 1 deletion.
11 changes: 11 additions & 0 deletions ibis/expr/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import functools
import numbers
import operator
from collections import Counter
from typing import TYPE_CHECKING, Any, NamedTuple, overload

import ibis.expr.builders as bl
Expand Down Expand Up @@ -476,6 +477,16 @@ def _memtable_from_dataframe(
(f"col{i:d}" for i in builtins.range(len(cols))),
)
df = df.rename(columns=dict(zip(cols, newcols)))

# verify that the DataFrame has no duplicate column names because ibis
# doesn't allow that
cols = df.columns
dupes = [name for name, count in Counter(cols).items() if count > 1]
if dupes:
raise IbisInputError(
f"Duplicate column names found in DataFrame when constructing memtable: {dupes}"
)

op = ops.InMemoryTable(
name=name if name is not None else util.gen_name("pandas_memtable"),
schema=sch.infer(df) if schema is None else schema,
Expand Down
9 changes: 8 additions & 1 deletion ibis/expr/tests/test_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
import ibis.expr.datatypes as dt
import ibis.expr.schema as sch
from ibis import _
from ibis.common.exceptions import IntegrityError
from ibis.common.exceptions import IbisInputError, IntegrityError


def test_schema_from_names_types():
Expand Down Expand Up @@ -117,3 +117,10 @@ def test_repr_deferred_with_exprs(f, sol):
expr = f(t)
res = repr(expr)
assert res == sol


def test_duplicate_columns_in_memtable_not_allowed():
df = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "a"])

with pytest.raises(IbisInputError, match="Duplicate column names"):
ibis.memtable(df)

0 comments on commit 4937b48

Please sign in to comment.