def infer_pd_series_spark_type(s: pd.Series) -> types.DataType:
"""Infer Spark DataType from pandas Series dtype.
:param s: :class:`pandas.Series` to be inferred
:return: the inferred Spark data type
"""
dt = s.dtype
if dt == np.dtype('object'):
if len(s) == 0 or s.isnull().all():
raise ValueError("can not infer schema from empty or null dataset")
return types.from_arrow_type(pa.Array.from_pandas(s).type)
elif is_datetime64_dtype(dt) or is_datetime64tz_dtype(dt):
return types.TimestampType()
else:
return types.from_arrow_type(pa.from_numpy_dtype(dt))
def infer_pd_series_spark_type(s: pd.Series) -> types.DataType:
"""Infer Spark DataType from pandas Series dtype.
:param s: :class:`pandas.Series` to be inferred
:return: the inferred Spark data type
"""
dt = s.dtype
if dt == np.dtype('object'):
if len(s) == 0 or s.isnull().all():
raise ValueError("can not infer schema from empty or null dataset")
return types.from_arrow_type(pa.Array.from_pandas(s).type)
elif is_datetime64_dtype(dt) or is_datetime64tz_dtype(dt):
return types.TimestampType()
else:
return types.from_arrow_type(pa.from_numpy_dtype(dt))
def arrow_to_pandas(self, arrow_column):
from pyspark.sql.types import from_arrow_type, \
_arrow_column_to_pandas, _check_series_localize_timestamps
s = _arrow_column_to_pandas(arrow_column, from_arrow_type(arrow_column.type))
s = _check_series_localize_timestamps(s, self._timezone)
return s
def arrow_to_pandas(self, arrow_column):
from pyspark.sql.types import from_arrow_type, \
_check_series_convert_date, _check_series_localize_timestamps
s = arrow_column.to_pandas()
s = _check_series_convert_date(s, from_arrow_type(arrow_column.type))
s = _check_series_localize_timestamps(s, self._timezone)
return s
def load_stream(self, stream):
"""
Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series.
"""
batches = super(ArrowStreamPandasSerializer, self).load_stream(stream)
import pyarrow as pa
from pyspark.sql.types import from_arrow_type
for batch in batches:
yield [self.arrow_to_pandas(c, from_arrow_type(c.type))
for c in pa.Table.from_batches([batch]).itercolumns()]
def load_stream(self, stream):
"""
Deserialize ArrowRecordBatches to an Arrow table and return as a list of pandas.Series.
"""
batches = super(ArrowStreamPandasSerializer, self).load_stream(stream)
import pyarrow as pa
from pyspark.sql.types import from_arrow_type
for batch in batches:
yield [
self.arrow_to_pandas(c, from_arrow_type(c.type))
for c in pa.Table.from_batches([batch]).itercolumns()
]
def _from_arrow_type(dt: pa.DataType) -> pt.DataType:
if is_struct(dt):
return pt.StructType([
pt.StructField(
# field.name, _from_arrow_type(field.type), nullable=field.nullable
field.name,
_from_arrow_type(field.type),
nullable=True,
) for field in dt
])
elif is_list(dt):
if is_timestamp(dt.value_type):
raise TypeError( # pragma: no cover
"Spark: unsupported type in conversion from Arrow: " + str(dt))
return pt.ArrayType(_from_arrow_type(dt.value_type))
return from_arrow_type(dt)
def infer_pd_series_spark_type(s: pd.Series) -> types.DataType:
"""Infer Spark DataType from pandas Series dtype.
:param s: :class:`pandas.Series` to be inferred
:return: the inferred Spark data type
"""
dt = s.dtype
if dt == np.dtype("object"):
if len(s) == 0 or s.isnull().all():
return types.NullType()
elif hasattr(s.iloc[0], "__UDT__"):
return s.iloc[0].__UDT__
else:
return from_arrow_type(pa.Array.from_pandas(s).type)
else:
return as_spark_type(dt)
def infer_pd_series_spark_type(s: pd.Series) -> types.DataType:
"""Infer Spark DataType from pandas Series dtype.
:param s: :class:`pandas.Series` to be inferred
:return: the inferred Spark data type
"""
dt = s.dtype
if dt == np.dtype("object"):
if len(s) == 0 or s.isnull().all():
raise ValueError("can not infer schema from empty or null dataset")
elif hasattr(s[0], "__UDT__"):
return s[0].__UDT__
else:
return from_arrow_type(pa.Array.from_pandas(s).type)
else:
return as_spark_type(dt)
def infer_pd_series_spark_type(pser: pd.Series, dtype: Dtype) -> types.DataType:
"""Infer Spark DataType from pandas Series dtype.
:param pser: :class:`pandas.Series` to be inferred
:param dtype: the Series' dtype
:return: the inferred Spark data type
"""
if dtype == np.dtype("object"):
if len(pser) == 0 or pser.isnull().all():
return types.NullType()
elif hasattr(pser.iloc[0], "__UDT__"):
return pser.iloc[0].__UDT__
else:
return from_arrow_type(pa.Array.from_pandas(pser).type)
elif isinstance(dtype, CategoricalDtype):
# `pser` must already be converted to codes.
return as_spark_type(pser.dtype)
else:
return as_spark_type(dtype)
def _create_from_pandas_with_arrow(self, pdf, schema, timezone):
"""
Create a DataFrame from a given pandas.DataFrame by slicing it into partitions, converting
to Arrow data, then sending to the JVM to parallelize. If a schema is passed in, the
data types will be used to coerce the data in Pandas to Arrow conversion.
"""
from distutils.version import LooseVersion
from pyspark.serializers import ArrowStreamPandasSerializer
from pyspark.sql.types import from_arrow_type, to_arrow_type, TimestampType
from pyspark.sql.utils import require_minimum_pandas_version, \
require_minimum_pyarrow_version
require_minimum_pandas_version()
require_minimum_pyarrow_version()
from pandas.api.types import is_datetime64_dtype, is_datetime64tz_dtype
import pyarrow as pa
# Create the Spark schema from list of names passed in with Arrow types
if isinstance(schema, (list, tuple)):
if LooseVersion(pa.__version__) < LooseVersion("0.12.0"):
temp_batch = pa.RecordBatch.from_pandas(pdf[0:100], preserve_index=False)
arrow_schema = temp_batch.schema
else:
arrow_schema = pa.Schema.from_pandas(pdf, preserve_index=False)
struct = StructType()
for name, field in zip(schema, arrow_schema):
struct.add(name, from_arrow_type(field.type), nullable=field.nullable)
schema = struct
# Determine arrow types to coerce data when creating batches
if isinstance(schema, StructType):
arrow_types = [to_arrow_type(f.dataType) for f in schema.fields]
elif isinstance(schema, DataType):
raise ValueError("Single data type %s is not supported with Arrow" % str(schema))
else:
# Any timestamps must be coerced to be compatible with Spark
arrow_types = [to_arrow_type(TimestampType())
if is_datetime64_dtype(t) or is_datetime64tz_dtype(t) else None
for t in pdf.dtypes]
# Slice the DataFrame to be batched
step = -(-len(pdf) // self.sparkContext.defaultParallelism) # round int up
pdf_slices = (pdf[start:start + step] for start in xrange(0, len(pdf), step))
# Create list of Arrow (columns, type) for serializer dump_stream
arrow_data = [[(c, t) for (_, c), t in zip(pdf_slice.iteritems(), arrow_types)]
for pdf_slice in pdf_slices]
jsqlContext = self._wrapped._jsqlContext
safecheck = self._wrapped._conf.arrowSafeTypeConversion()
col_by_name = True # col by name only applies to StructType columns, can't happen here
ser = ArrowStreamPandasSerializer(timezone, safecheck, col_by_name)
def reader_func(temp_filename):
return self._jvm.PythonSQLUtils.readArrowStreamFromFile(jsqlContext, temp_filename)
def create_RDD_server():
return self._jvm.ArrowRDDServer(jsqlContext)
# Create Spark DataFrame from Arrow stream file, using one batch per partition
jrdd = self._sc._serialize_to_jvm(arrow_data, ser, reader_func, create_RDD_server)
jdf = self._jvm.PythonSQLUtils.toDataFrame(jrdd, schema.json(), jsqlContext)
df = DataFrame(jdf, self._wrapped)
df._schema = schema
return df
def pyarrow_to_spark_data_type(dtype):
# PySpark will interpret list types as Arrays, but for ML applications we want to default to
# treating these as DenseVectors.
if pa.types.is_list(dtype):
return DenseVector
return type(from_arrow_type(dtype))
