1from __future__ import annotations
2
3from collections import defaultdict
4from collections.abc import (
5 Hashable,
6 Iterable,
7)
8import itertools
9from typing import TYPE_CHECKING
10
11import numpy as np
12
13from pandas._libs import missing as libmissing
14from pandas._libs.sparse import IntIndex
15from pandas.util._decorators import set_module
16
17from pandas.core.dtypes.common import (
18 is_integer_dtype,
19 is_list_like,
20 is_object_dtype,
21 pandas_dtype,
22)
23from pandas.core.dtypes.dtypes import (
24 ArrowDtype,
25 CategoricalDtype,
26)
27
28from pandas.core.arrays import SparseArray
29from pandas.core.arrays.categorical import factorize_from_iterable
30from pandas.core.arrays.string_ import StringDtype
31from pandas.core.frame import DataFrame
32from pandas.core.indexes.api import (
33 Index,
34 default_index,
35)
36from pandas.core.series import Series
37
38if TYPE_CHECKING:
39 from pandas._typing import NpDtype
40
41
42@set_module("pandas")
43def get_dummies(
44 data,
45 prefix=None,
46 prefix_sep: str | Iterable[str] | dict[str, str] = "_",
47 dummy_na: bool = False,
48 columns=None,
49 sparse: bool = False,
50 drop_first: bool = False,
51 dtype: NpDtype | None = None,
52) -> DataFrame:
53 """
54 Convert categorical variable into dummy/indicator variables.
55
56 Each variable is converted in as many 0/1 variables as there are different
57 values. Columns in the output are each named after a value; if the input is
58 a DataFrame, the name of the original variable is prepended to the value.
59
60 Parameters
61 ----------
62 data : array-like, Series, or DataFrame
63 Data of which to get dummy indicators.
64 prefix : str, list of str, or dict of str, default None
65 A string to be prepended to DataFrame column names.
66 Pass a list with length equal to the number of columns
67 when calling get_dummies on a DataFrame. Alternatively, `prefix`
68 can be a dictionary mapping column names to prefixes.
69 prefix_sep : str, list of str, or dict of str, default '_'
70 Should you choose to prepend DataFrame column names with a prefix, this
71 is the separator/delimiter to use between the two. Alternatively,
72 `prefix_sep` can be a list with length equal to the number of columns,
73 or a dictionary mapping column names to separators.
74 dummy_na : bool, default False
75 If True, a NaN indicator column will be added even if no NaN values are present.
76 If False, NA values are encoded as all zero.
77 columns : list-like, default None
78 Column names in the DataFrame to be encoded.
79 If `columns` is None then all the columns with
80 `object`, `string`, or `category` dtype will be converted.
81 sparse : bool, default False
82 Whether the dummy-encoded columns should be backed by
83 a :class:`SparseArray` (True) or a regular NumPy array (False).
84 drop_first : bool, default False
85 Whether to get k-1 dummies out of k categorical levels by removing the
86 first level.
87 dtype : dtype, default bool
88 Data type for new columns. Only a single dtype is allowed.
89
90 Returns
91 -------
92 DataFrame
93 Dummy-coded data. If `data` contains other columns than the
94 dummy-coded one(s), these will be prepended, unaltered, to the result.
95
96 See Also
97 --------
98 Series.str.get_dummies : Convert Series of strings to dummy codes.
99 :func:`~pandas.from_dummies` : Convert dummy codes to categorical ``DataFrame``.
100
101 Notes
102 -----
103 Reference :ref:`the user guide <reshaping.dummies>` for more examples.
104
105 Examples
106 --------
107 >>> s = pd.Series(list("abca"))
108
109 >>> pd.get_dummies(s)
110 a b c
111 0 True False False
112 1 False True False
113 2 False False True
114 3 True False False
115
116 >>> s1 = ["a", "b", np.nan]
117
118 >>> pd.get_dummies(s1)
119 a b
120 0 True False
121 1 False True
122 2 False False
123
124 >>> pd.get_dummies(s1, dummy_na=True)
125 a b NaN
126 0 True False False
127 1 False True False
128 2 False False True
129
130 >>> df = pd.DataFrame({"A": ["a", "b", "a"], "B": ["b", "a", "c"], "C": [1, 2, 3]})
131
132 >>> pd.get_dummies(df, prefix=["col1", "col2"])
133 C col1_a col1_b col2_a col2_b col2_c
134 0 1 True False False True False
135 1 2 False True True False False
136 2 3 True False False False True
137
138 >>> pd.get_dummies(pd.Series(list("abcaa")))
139 a b c
140 0 True False False
141 1 False True False
142 2 False False True
143 3 True False False
144 4 True False False
145
146 >>> pd.get_dummies(pd.Series(list("abcaa")), drop_first=True)
147 b c
148 0 False False
149 1 True False
150 2 False True
151 3 False False
152 4 False False
153
154 >>> pd.get_dummies(pd.Series(list("abc")), dtype=float)
155 a b c
156 0 1.0 0.0 0.0
157 1 0.0 1.0 0.0
158 2 0.0 0.0 1.0
159 """
160 from pandas.core.reshape.concat import concat
161
162 dtypes_to_encode = ["object", "string", "category"]
163
164 if isinstance(data, DataFrame):
165 # determine columns being encoded
166 if columns is None:
167 data_to_encode = data.select_dtypes(include=dtypes_to_encode)
168 elif not is_list_like(columns):
169 raise TypeError("Input must be a list-like for parameter `columns`")
170 else:
171 data_to_encode = data[columns]
172
173 # validate prefixes and separator to avoid silently dropping cols
174 def check_len(item, name: str) -> None:
175 if is_list_like(item):
176 if not len(item) == data_to_encode.shape[1]:
177 len_msg = (
178 f"Length of '{name}' ({len(item)}) did not match the "
179 "length of the columns being encoded "
180 f"({data_to_encode.shape[1]})."
181 )
182 raise ValueError(len_msg)
183
184 check_len(prefix, "prefix")
185 check_len(prefix_sep, "prefix_sep")
186
187 if isinstance(prefix, str):
188 prefix = itertools.repeat(prefix, len(data_to_encode.columns))
189 if isinstance(prefix, dict):
190 prefix = [prefix[col] for col in data_to_encode.columns]
191
192 if prefix is None:
193 prefix = data_to_encode.columns
194
195 # validate separators
196 if isinstance(prefix_sep, str):
197 prefix_sep = itertools.repeat(prefix_sep, len(data_to_encode.columns))
198 elif isinstance(prefix_sep, dict):
199 prefix_sep = [prefix_sep[col] for col in data_to_encode.columns]
200
201 with_dummies: list[DataFrame]
202 if data_to_encode.shape == data.shape:
203 # Encoding the entire df, do not prepend any dropped columns
204 with_dummies = []
205 elif columns is not None:
206 # Encoding only cols specified in columns. Get all cols not in
207 # columns to prepend to result.
208 with_dummies = [data.drop(columns, axis=1)]
209 else:
210 # Encoding only object and category dtype columns. Get remaining
211 # columns to prepend to result.
212 with_dummies = [data.select_dtypes(exclude=dtypes_to_encode)]
213
214 for col, pre, sep in zip(
215 data_to_encode.items(), prefix, prefix_sep, strict=True
216 ):
217 # col is (column_name, column), use just column data here
218 dummy = _get_dummies_1d(
219 col[1],
220 prefix=pre,
221 prefix_sep=sep,
222 dummy_na=dummy_na,
223 sparse=sparse,
224 drop_first=drop_first,
225 dtype=dtype,
226 )
227 with_dummies.append(dummy)
228 result = concat(with_dummies, axis=1)
229 else:
230 result = _get_dummies_1d(
231 data,
232 prefix,
233 prefix_sep,
234 dummy_na,
235 sparse=sparse,
236 drop_first=drop_first,
237 dtype=dtype,
238 )
239 return result
240
241
242def _get_dummies_1d(
243 data,
244 prefix,
245 prefix_sep: str | Iterable[str] | dict[str, str] = "_",
246 dummy_na: bool = False,
247 sparse: bool = False,
248 drop_first: bool = False,
249 dtype: NpDtype | None = None,
250) -> DataFrame:
251 from pandas.core.reshape.concat import concat
252
253 # Series avoids inconsistent NaN handling
254 codes, levels = factorize_from_iterable(Series(data, copy=False))
255
256 if dtype is None and hasattr(data, "dtype"):
257 input_dtype = data.dtype
258 if isinstance(input_dtype, CategoricalDtype):
259 input_dtype = input_dtype.categories.dtype
260
261 if isinstance(input_dtype, ArrowDtype):
262 import pyarrow as pa
263
264 dtype = ArrowDtype(pa.bool_()) # type: ignore[assignment]
265 elif (
266 isinstance(input_dtype, StringDtype)
267 and input_dtype.na_value is libmissing.NA
268 ):
269 dtype = pandas_dtype("boolean") # type: ignore[assignment]
270 else:
271 dtype = np.dtype(bool)
272 elif dtype is None:
273 dtype = np.dtype(bool)
274
275 _dtype = pandas_dtype(dtype)
276
277 if is_object_dtype(_dtype):
278 raise ValueError("dtype=object is not a valid dtype for get_dummies")
279
280 def get_empty_frame(data) -> DataFrame:
281 index: Index | np.ndarray
282 if isinstance(data, Series):
283 index = data.index
284 else:
285 index = default_index(len(data))
286 return DataFrame(index=index)
287
288 # if all NaN
289 if not dummy_na and len(levels) == 0:
290 return get_empty_frame(data)
291
292 codes = codes.copy()
293 if dummy_na:
294 codes[codes == -1] = len(levels)
295 levels = levels.insert(len(levels), np.nan)
296
297 # if dummy_na, we just fake a nan level. drop_first will drop it again
298 if drop_first and len(levels) == 1:
299 return get_empty_frame(data)
300
301 number_of_cols = len(levels)
302
303 if prefix is None:
304 dummy_cols = levels
305 else:
306 dummy_cols = Index([f"{prefix}{prefix_sep}{level}" for level in levels])
307
308 index: Index | None
309 if isinstance(data, Series):
310 index = data.index
311 else:
312 index = None
313
314 if sparse:
315 fill_value: bool | float
316 if is_integer_dtype(dtype):
317 fill_value = 0
318 elif dtype == np.dtype(bool):
319 fill_value = False
320 else:
321 fill_value = 0.0
322
323 sparse_series = []
324 N = len(data)
325 sp_indices: list[list] = [[] for _ in range(len(dummy_cols))]
326 mask = codes != -1
327 codes = codes[mask]
328 n_idx = np.arange(N)[mask]
329
330 for ndx, code in zip(n_idx, codes, strict=True):
331 sp_indices[code].append(ndx)
332
333 if drop_first:
334 # remove first categorical level to avoid perfect collinearity
335 # GH12042
336 sp_indices = sp_indices[1:]
337 dummy_cols = dummy_cols[1:]
338 for col, ixs in zip(dummy_cols, sp_indices, strict=True):
339 sarr = SparseArray(
340 np.ones(len(ixs), dtype=dtype),
341 sparse_index=IntIndex(N, ixs),
342 fill_value=fill_value,
343 dtype=dtype,
344 )
345 sparse_series.append(Series(data=sarr, index=index, name=col, copy=False))
346
347 return concat(sparse_series, axis=1)
348
349 else:
350 # ensure ndarray layout is column-major
351 shape = len(codes), number_of_cols
352 dummy_dtype: NpDtype
353 if isinstance(_dtype, np.dtype):
354 dummy_dtype = _dtype
355 else:
356 dummy_dtype = np.bool_
357 dummy_mat = np.zeros(shape=shape, dtype=dummy_dtype, order="F")
358 dummy_mat[np.arange(len(codes)), codes] = 1
359
360 if not dummy_na:
361 # reset NaN GH4446
362 dummy_mat[codes == -1] = 0
363
364 if drop_first:
365 # remove first GH12042
366 dummy_mat = dummy_mat[:, 1:]
367 dummy_cols = dummy_cols[1:]
368 return DataFrame(dummy_mat, index=index, columns=dummy_cols, dtype=_dtype)
369
370
371@set_module("pandas")
372def from_dummies(
373 data: DataFrame,
374 sep: None | str = None,
375 default_category: None | Hashable | dict[str, Hashable] = None,
376) -> DataFrame:
377 """
378 Create a categorical ``DataFrame`` from a ``DataFrame`` of dummy variables.
379
380 Inverts the operation performed by :func:`~pandas.get_dummies`.
381
382 Parameters
383 ----------
384 data : DataFrame
385 Data which contains dummy-coded variables in form of integer columns of
386 1's and 0's.
387 sep : str, default None
388 Separator used in the column names of the dummy categories they are
389 character indicating the separation of the categorical names from the prefixes.
390 For example, if your column names are 'prefix_A' and 'prefix_B',
391 you can strip the underscore by specifying sep='_'.
392 default_category : None, Hashable or dict of Hashables, default None
393 The default category is the implied category when a value has none of the
394 listed categories specified with a one, i.e. if all dummies in a row are
395 zero. Can be a single value for all variables or a dict directly mapping
396 the default categories to a prefix of a variable. The default category
397 will be coerced to the dtype of ``data.columns`` if such coercion is
398 lossless, and will raise otherwise.
399
400 Returns
401 -------
402 DataFrame
403 Categorical data decoded from the dummy input-data.
404
405 Raises
406 ------
407 ValueError
408 * When the input ``DataFrame`` ``data`` contains NA values.
409 * When the input ``DataFrame`` ``data`` contains column names with separators
410 that do not match the separator specified with ``sep``.
411 * When a ``dict`` passed to ``default_category`` does not include an implied
412 category for each prefix.
413 * When a value in ``data`` has more than one category assigned to it.
414 * When ``default_category=None`` and a value in ``data`` has no category
415 assigned to it.
416 TypeError
417 * When the input ``data`` is not of type ``DataFrame``.
418 * When the input ``DataFrame`` ``data`` contains non-dummy data.
419 * When the passed ``sep`` is of a wrong data type.
420 * When the passed ``default_category`` is of a wrong data type.
421
422 See Also
423 --------
424 :func:`~pandas.get_dummies` : Convert ``Series`` or ``DataFrame`` to dummy codes.
425 :class:`~pandas.Categorical` : Represent a categorical variable in classic.
426
427 Notes
428 -----
429 The columns of the passed dummy data should only include 1's and 0's,
430 or boolean values.
431
432 Examples
433 --------
434 >>> df = pd.DataFrame({"a": [1, 0, 0, 1], "b": [0, 1, 0, 0], "c": [0, 0, 1, 0]})
435
436 >>> df
437 a b c
438 0 1 0 0
439 1 0 1 0
440 2 0 0 1
441 3 1 0 0
442
443 >>> pd.from_dummies(df)
444 0 a
445 1 b
446 2 c
447 3 a
448
449 >>> df = pd.DataFrame(
450 ... {
451 ... "col1_a": [1, 0, 1],
452 ... "col1_b": [0, 1, 0],
453 ... "col2_a": [0, 1, 0],
454 ... "col2_b": [1, 0, 0],
455 ... "col2_c": [0, 0, 1],
456 ... }
457 ... )
458
459 >>> df
460 col1_a col1_b col2_a col2_b col2_c
461 0 1 0 0 1 0
462 1 0 1 1 0 0
463 2 1 0 0 0 1
464
465 >>> pd.from_dummies(df, sep="_")
466 col1 col2
467 0 a b
468 1 b a
469 2 a c
470
471 >>> df = pd.DataFrame(
472 ... {
473 ... "col1_a": [1, 0, 0],
474 ... "col1_b": [0, 1, 0],
475 ... "col2_a": [0, 1, 0],
476 ... "col2_b": [1, 0, 0],
477 ... "col2_c": [0, 0, 0],
478 ... }
479 ... )
480
481 >>> df
482 col1_a col1_b col2_a col2_b col2_c
483 0 1 0 0 1 0
484 1 0 1 1 0 0
485 2 0 0 0 0 0
486
487 >>> pd.from_dummies(df, sep="_", default_category={"col1": "d", "col2": "e"})
488 col1 col2
489 0 a b
490 1 b a
491 2 d e
492 """
493 from pandas.core.reshape.concat import concat
494
495 if not isinstance(data, DataFrame):
496 raise TypeError(
497 "Expected 'data' to be a 'DataFrame'; "
498 f"Received 'data' of type: {type(data).__name__}"
499 )
500
501 col_isna_mask = data.isna().any()
502
503 if col_isna_mask.any():
504 raise ValueError(
505 f"Dummy DataFrame contains NA value in column: '{col_isna_mask.idxmax()}'"
506 )
507
508 # index data with a list of all columns that are dummies
509 try:
510 data_to_decode = data.astype("boolean")
511 except TypeError as err:
512 raise TypeError("Passed DataFrame contains non-dummy data") from err
513
514 # collect prefixes and get lists to slice data for each prefix
515 variables_slice = defaultdict(list)
516 if sep is None:
517 variables_slice[""] = list(data.columns)
518 elif isinstance(sep, str):
519 for col in data_to_decode.columns:
520 prefix = col.split(sep)[0]
521 if len(prefix) == len(col):
522 raise ValueError(f"Separator not specified for column: {col}")
523 variables_slice[prefix].append(col)
524 else:
525 raise TypeError(
526 "Expected 'sep' to be of type 'str' or 'None'; "
527 f"Received 'sep' of type: {type(sep).__name__}"
528 )
529
530 if default_category is not None:
531 if isinstance(default_category, dict):
532 if not len(default_category) == len(variables_slice):
533 len_msg = (
534 f"Length of 'default_category' ({len(default_category)}) "
535 f"did not match the length of the columns being encoded "
536 f"({len(variables_slice)})"
537 )
538 raise ValueError(len_msg)
539 elif isinstance(default_category, Hashable):
540 default_category = dict(
541 zip(
542 variables_slice,
543 [default_category] * len(variables_slice),
544 strict=True,
545 )
546 )
547 else:
548 raise TypeError(
549 "Expected 'default_category' to be of type "
550 "'None', 'Hashable', or 'dict'; "
551 "Received 'default_category' of type: "
552 f"{type(default_category).__name__}"
553 )
554
555 cat_data = {}
556 for prefix, prefix_slice in variables_slice.items():
557 if sep is None:
558 cats = prefix_slice.copy()
559 else:
560 cats = [col[len(prefix + sep) :] for col in prefix_slice]
561 assigned = data_to_decode.loc[:, prefix_slice].sum(axis=1)
562 if any(assigned > 1):
563 raise ValueError(
564 "Dummy DataFrame contains multi-assignment(s); "
565 f"First instance in row: {assigned.idxmax()}"
566 )
567 if any(assigned == 0):
568 if isinstance(default_category, dict):
569 cats.append(default_category[prefix])
570 else:
571 raise ValueError(
572 "Dummy DataFrame contains unassigned value(s); "
573 f"First instance in row: {assigned.idxmin()}"
574 )
575 data_slice = concat(
576 (data_to_decode.loc[:, prefix_slice], assigned == 0), axis=1
577 )
578 else:
579 data_slice = data_to_decode.loc[:, prefix_slice]
580 cats_array = data._constructor_sliced(cats, dtype=data.columns.dtype)
581 # get indices of True entries along axis=1
582 true_values = data_slice.idxmax(axis=1)
583 indexer = data_slice.columns.get_indexer_for(true_values)
584 cat_data[prefix] = cats_array.take(indexer).set_axis(data.index)
585
586 result = DataFrame(cat_data)
587 if sep is not None:
588 result.columns = result.columns.astype(data.columns.dtype)
589 return result