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