|
1 | 1 | from typing import List, Tuple, Iterator, Iterable, Any, Optional, Union |
| 2 | +from collections.abc import Sized |
2 | 3 | from functools import wraps |
3 | 4 |
|
4 | 5 | import pandas as pd |
@@ -133,24 +134,30 @@ def emit(self, *args) -> None: |
133 | 134 | self._current_context.emit(*args) |
134 | 135 |
|
135 | 136 |
|
136 | | -def get_scalar_input(inp: Any) -> Iterable[Tuple[Any, ...]]: |
| 137 | +def get_scalar_input(inp: Any) -> Iterable[Iterable[Any]]: |
137 | 138 | """ |
138 | 139 | Figures out if the SCALAR parameters are provided as a scalar value or a tuple |
139 | 140 | and also if there is a wrapping container around. |
140 | | - Unless the parameters are already in a wrapping container returns parameters as a tuple wrapped |
141 | | - into a one-item list, e.g [(param1[, param2, ...)]. Otherwise, returns the original input. |
| 141 | + Unless the parameters are already in a wrapping Sized container, returns parameters as an iterable |
| 142 | + wrapped into a one-item list, e.g [(param1, [param2, ...])]. Otherwise, returns the original input. |
142 | 143 |
|
143 | 144 | :param inp: Input parameters. |
144 | 145 | """ |
145 | 146 |
|
146 | | - if isinstance(inp, Iterable) and not isinstance(inp, str): |
147 | | - row1 = next(iter(inp)) |
148 | | - if isinstance(row1, Iterable) and not isinstance(row1, str): |
149 | | - return inp |
150 | | - else: |
151 | | - return [inp] |
152 | | - else: |
153 | | - return [(inp,)] |
| 147 | + if inp is not None: |
| 148 | + if (not isinstance(inp, Iterable)) or isinstance(inp, str): |
| 149 | + return [(inp,)] |
| 150 | + try: |
| 151 | + row1 = next(iter(inp)) |
| 152 | + if (not isinstance(row1, Iterable)) or isinstance(row1, str): |
| 153 | + return [inp] |
| 154 | + elif not isinstance(inp, Sized): |
| 155 | + return list(inp) |
| 156 | + else: |
| 157 | + return inp |
| 158 | + except StopIteration: |
| 159 | + pass |
| 160 | + return [tuple()] |
154 | 161 |
|
155 | 162 |
|
156 | 163 | class StandaloneMockContext(UDFContext): |
@@ -205,20 +212,18 @@ def get_dataframe(self, num_rows='all', start_col=0): |
205 | 212 | columns_ = [column.name for column in self._metadata.input_columns[start_col:]] |
206 | 213 |
|
207 | 214 | i = 0 |
208 | | - df = None |
| 215 | + dfs: list[pd.DataFrame] = [] |
209 | 216 | while num_rows == 'all' or i < num_rows: |
210 | | - df_current = pd.DataFrame.from_records( |
211 | | - [self._data[start_col:]], columns=columns_) |
212 | | - if df is None: |
213 | | - df = df_current |
214 | | - else: |
215 | | - df = df.append(df_current) |
| 217 | + dfs.append(pd.DataFrame.from_records( |
| 218 | + [self._data[start_col:]], columns=columns_)) |
216 | 219 | if not self.next(): |
217 | 220 | break |
218 | 221 | i += 1 |
219 | | - if df is not None: |
220 | | - df = df.reset_index(drop=True) |
221 | | - return df |
| 222 | + if dfs: |
| 223 | + df = pd.concat(dfs, ignore_index=True) |
| 224 | + df.reset_index(inplace=True, drop=True) |
| 225 | + return df |
| 226 | + return None |
222 | 227 |
|
223 | 228 | def __getattr__(self, name): |
224 | 229 | return None if self._data is None else self._data[self._name_position_map[name]] |
|
0 commit comments