|
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