dchvn commented on a change in pull request #34181:
URL: https://github.com/apache/spark/pull/34181#discussion_r722083275
##########
File path: python/pyspark/pandas/typedef/typehints.py
##########
@@ -726,131 +729,108 @@ def create_tuple_for_frame_type(params: Any) -> object:
... # doctest: +ELLIPSIS
typing.Tuple[...IndexNameType, ...NameType]
"""
- return Tuple[_extract_types(params)]
+ return Tuple[_to_type_holders(params)]
-def _extract_types(params: Any) -> Tuple:
- origin = params
+def _to_type_holders(params: Any) -> Tuple:
+ from pyspark.pandas.typedef import NameTypeHolder, IndexNameTypeHolder
- params = _to_tuple_of_params(params)
+ is_with_index = (
+ isinstance(params, tuple)
+ and len(params) == 2
+ and isinstance(params[1], (zip, list, pd.Series))
+ )
- if _is_named_params(params):
- # Example:
- # DataFrame["id": int, "A": int]
- new_params = _address_named_type_hoders(params, is_index=False)
- return tuple(new_params)
- elif len(params) == 2 and isinstance(params[1], (zip, list, pd.Series)):
- # Example:
- # DataFrame[int, [int, int]]
- # DataFrame[pdf.index.dtype, pdf.dtypes]
- # DataFrame[("index", int), [("id", int), ("A", int)]]
- # DataFrame[(pdf.index.name, pdf.index.dtype), zip(pdf.columns,
pdf.dtypes)]
- #
- # DataFrame[[int, int], [int, int]]
- # DataFrame[pdf.index.dtypes, pdf.dtypes]
- # DataFrame[[("index", int), ("index-2", int)], [("id", int), ("A",
int)]]
- # DataFrame[zip(pdf.index.names, pdf.index.dtypes), zip(pdf.columns,
pdf.dtypes)]
+ if is_with_index:
+ # With index
+ # DataFrame[index_type, [type, ...]]
+ # DataFrame[dtype instance, dtypes instance]
+ # DataFrame[[index_type, ...], [type, ...]]
+ # DataFrame[dtypes instance, dtypes instance]
+ # DataFrame[(index_name, index_type), [(name, type), ...]]
+ # DataFrame[(index_name, index_type), zip(names, types)]
+ # DataFrame[[(index_name, index_type), ...], [(name, type), ...]]
+ # DataFrame[zip(index_names, index_types), zip(names, types)]
+ def is_list_of_pairs(p: Any) -> bool:
+ return (
+ isinstance(p, list) and len(p) >= 1 and isinstance(p[0],
tuple) and len(p[0]) == 2
+ )
Review comment:
Should we check all elements?
``` python
return (
isinstance(p, list)
and (len(p) >= 1)
and all(isinstance(param, tuple) and (len(param) == 2) for
param in p)
)
```
for these cases:
```
>>> ps.DataFrame[int, [("index", int), ["index-2", int]]]
typing.Tuple[pyspark.pandas.typedef.typehints.IndexNameType,
pyspark.pandas.typedef.typehints.NameType,
pyspark.pandas.typedef.typehints.NameType]
>>> ps.DataFrame[int, [("index", int), ("index-2", int, None)]]
typing.Tuple[pyspark.pandas.typedef.typehints.IndexNameType,
pyspark.pandas.typedef.typehints.NameType,
pyspark.pandas.typedef.typehints.NameType]
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
