def infer_pd_series_spark_type( pser: pd.Series, dtype: Dtype, prefer_timestamp_ntz: bool = False) -> types.DataType: """Infer Spark DataType from pandas Series dtype. :param pser: :class:`pandas.Series` to be inferred :param dtype: the Series' dtype :param prefer_timestamp_ntz: if true, infers datetime without timezone as TimestampNTZType type. If false, infers it as TimestampType. :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, prefer_timestamp_ntz) elif isinstance(dtype, CategoricalDtype): if isinstance(pser.dtype, CategoricalDtype): return as_spark_type(pser.cat.codes.dtype, prefer_timestamp_ntz=prefer_timestamp_ntz) else: # `pser` must already be converted to codes. return as_spark_type(pser.dtype, prefer_timestamp_ntz=prefer_timestamp_ntz) else: return as_spark_type(dtype, prefer_timestamp_ntz=prefer_timestamp_ntz)
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 ray_dataset_to_spark_dataframe(spark: sql.SparkSession, arrow_schema: "pa.lib.Schema", blocks: List[ObjectRef], locations: List[bytes]) -> DataFrame: if not isinstance(arrow_schema, pa.lib.Schema): raise RuntimeError(f"Schema is {type(arrow_schema)}, required pyarrow.lib.Schema. \n" \ f"to_spark does not support converting non-arrow ray datasets.") schema = StructType() for field in arrow_schema: schema.add(field.name, from_arrow_type(field.type), nullable=field.nullable) #TODO how to branch on type of block? sample = ray.get(blocks[0]) if isinstance(sample, bytes): return _convert_by_rdd(spark, blocks, locations, schema) elif isinstance(sample, pa.Table): return _convert_by_udf(spark, blocks, locations, schema) else: raise RuntimeError("ray.to_spark only supports arrow type blocks")
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 pyspark.sql import SparkSession from pyspark.sql.dataframe import DataFrame assert isinstance(self, SparkSession) from pyspark.sql.pandas.serializers import ArrowStreamPandasSerializer from pyspark.sql.types import TimestampType from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type from pyspark.sql.pandas.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)): arrow_schema = pa.Schema.from_pandas(pdf, preserve_index=False) struct = StructType() prefer_timestamp_ntz = is_timestamp_ntz_preferred() for name, field in zip(schema, arrow_schema): struct.add(name, from_arrow_type(field.type, prefer_timestamp_ntz), 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.iloc[start:start + step] for start in range(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))