def test_to_csv_with_path(self):
pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]})
kdf = ps.DataFrame(pdf)
tmp_dir = "{}/tmp1".format(self.tmp_dir)
kdf.to_csv(tmp_dir, num_files=1)
self._check_output(tmp_dir, pdf.to_csv(index=False))
tmp_dir = "{}/tmp2".format(self.tmp_dir)
self.assertRaises(KeyError, lambda: kdf.to_csv(tmp_dir, columns=["c"], num_files=1))
# non-string names
pdf = pd.DataFrame({10: [1, 2, 3], 20: ["a", "b", "c"]})
kdf = ps.DataFrame(pdf)
tmp_dir = "{}/tmp3".format(self.tmp_dir)
kdf.to_csv(tmp_dir, num_files=1)
self._check_output(tmp_dir, pdf.to_csv(index=False))
tmp_dir = "{}/tmp4".format(self.tmp_dir)
kdf.to_csv(tmp_dir, columns=[10], num_files=1)
self._check_output(tmp_dir, pdf.to_csv(columns=[10], index=False))
tmp_dir = "{}/tmp5".format(self.tmp_dir)
self.assertRaises(TypeError, lambda: kdf.to_csv(tmp_dir, columns=10, num_files=1))
def test_coalesce(self):
num_partitions = 10
psdf = ps.DataFrame({"age": [5, 5, 2, 2], "name": ["Bob", "Bob", "Alice", "Alice"]})
psdf = psdf.spark.repartition(num_partitions)
num_partitions -= 1
new_psdf = psdf.spark.coalesce(num_partitions)
self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
# Reserves Index
psdf = psdf.set_index("age")
num_partitions -= 1
new_psdf = psdf.spark.coalesce(num_partitions)
self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
# Reflects internal changes
psdf = psdf.reset_index()
psdf = psdf.set_index("name")
psdf2 = psdf + 1
num_partitions -= 1
self.assert_eq(psdf2.sort_index(), (psdf + 1).spark.coalesce(num_partitions).sort_index())
# Reserves MultiIndex
psdf = ps.DataFrame({"a": ["a", "b", "c"]}, index=[[1, 2, 3], [4, 5, 6]])
num_partitions -= 1
psdf = psdf.spark.repartition(num_partitions)
num_partitions -= 1
new_psdf = psdf.spark.coalesce(num_partitions)
self.assertEqual(new_psdf.to_spark().rdd.getNumPartitions(), num_partitions)
self.assert_eq(psdf.sort_index(), new_psdf.sort_index())
def test_expanding_count(self):
# The behaviour of Expanding.count are different between pandas>=1.0.0 and lower,
# and we're following the behaviour of latest version of pandas.
if LooseVersion(pd.__version__) >= LooseVersion("1.0.0"):
self._test_expanding_func("count")
else:
# Series
idx = np.random.rand(3)
psser = ps.Series([1, 2, 3], index=idx, name="a")
expected_result = pd.Series([None, 2.0, 3.0], index=idx, name="a")
self.assert_eq(psser.expanding(2).count().sort_index(), expected_result.sort_index())
self.assert_eq(psser.expanding(2).count().sum(), expected_result.sum())
# MultiIndex
midx = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y"), ("b", "z")])
psser = ps.Series([1, 2, 3], index=midx, name="a")
expected_result = pd.Series([None, 2.0, 3.0], index=midx, name="a")
self.assert_eq(psser.expanding(2).count().sort_index(), expected_result.sort_index())
# DataFrame
psdf = ps.DataFrame({"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0]})
expected_result = pd.DataFrame({"a": [None, 2.0, 3.0, 4.0], "b": [None, 2.0, 3.0, 4.0]})
self.assert_eq(psdf.expanding(2).count().sort_index(), expected_result.sort_index())
self.assert_eq(psdf.expanding(2).count().sum(), expected_result.sum())
# MultiIndex columns
idx = np.random.rand(4)
psdf = ps.DataFrame({"a": [1, 2, 3, 2], "b": [4.0, 2.0, 3.0, 1.0]}, index=idx)
psdf.columns = pd.MultiIndex.from_tuples([("a", "x"), ("a", "y")])
expected_result = pd.DataFrame(
{("a", "x"): [None, 2.0, 3.0, 4.0], ("a", "y"): [None, 2.0, 3.0, 4.0]},
index=idx,
)
self.assert_eq(psdf.expanding(2).count().sort_index(), expected_result.sort_index())
def test_to_csv(self):
pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": [4, 5, 6]}, index=[0, 1, 3])
kdf = ps.DataFrame(pdf)
self.assert_eq(kdf.to_csv(), pdf.to_csv(index=False))
self.assert_eq(kdf.to_csv(columns=["aa"]), pdf.to_csv(columns=["aa"], index=False))
self.assert_eq(kdf.aa.to_csv(), pdf.aa.to_csv(index=False, header=True))
pdf = pd.DataFrame({"a": [1, np.nan, 3], "b": ["one", "two", None]}, index=[0, 1, 3])
kdf = ps.from_pandas(pdf)
self.assert_eq(kdf.to_csv(na_rep="null"), pdf.to_csv(na_rep="null", index=False))
self.assert_eq(
kdf.a.to_csv(na_rep="null"), pdf.a.to_csv(na_rep="null", index=False, header=True)
)
self.assertRaises(KeyError, lambda: kdf.to_csv(columns=["ab"]))
pdf = pd.DataFrame({"a": [1.0, 2.0, 3.0], "b": [4.0, 5.0, 6.0]}, index=[0, 1, 3])
kdf = ps.from_pandas(pdf)
self.assert_eq(kdf.to_csv(), pdf.to_csv(index=False))
self.assert_eq(kdf.to_csv(header=False), pdf.to_csv(header=False, index=False))
self.assert_eq(kdf.to_csv(), pdf.to_csv(index=False))
# non-string names
pdf = pd.DataFrame({10: [1, 2, 3], 20: [4, 5, 6]}, index=[0, 1, 3])
kdf = ps.DataFrame(pdf)
self.assert_eq(kdf.to_csv(), pdf.to_csv(index=False))
self.assert_eq(kdf.to_csv(columns=[10]), pdf.to_csv(columns=[10], index=False))
self.assertRaises(TypeError, lambda: kdf.to_csv(columns=10))
def test_sql_with_index_col(self):
import pandas as pd
# Index
psdf = ps.DataFrame({
"A": [1, 2, 3],
"B": [4, 5, 6]
},
index=pd.Index(["a", "b", "c"], name="index"))
psdf_reset_index = psdf.reset_index()
actual = ps.sql("select * from {psdf_reset_index} where A > 1",
index_col="index")
expected = psdf.iloc[[1, 2]]
self.assert_eq(actual, expected)
# MultiIndex
psdf = ps.DataFrame(
{
"A": [1, 2, 3],
"B": [4, 5, 6]
},
index=pd.MultiIndex.from_tuples([("a", "b"), ("c", "d"),
("e", "f")],
names=["index1", "index2"]),
)
psdf_reset_index = psdf.reset_index()
actual = ps.sql("select * from {psdf_reset_index} where A > 1",
index_col=["index1", "index2"])
expected = psdf.iloc[[1, 2]]
self.assert_eq(actual, expected)
def test_sql_with_python_objects(self):
self.assert_eq(ps.sql("SELECT {col} as a FROM range(1)", col="lit"),
ps.DataFrame({"a": ["lit"]}))
self.assert_eq(
ps.sql("SELECT id FROM range(10) WHERE id IN {pred}",
col="lit",
pred=(1, 2, 3)),
ps.DataFrame({"id": [1, 2, 3]}),
)
def test_np_unsupported_frame(self):
psdf = self.psdf
with self.assertRaisesRegex(NotImplementedError,
"on-Spark.*not.*support.*sqrt.*"):
np.sqrt(psdf, psdf)
psdf1 = ps.DataFrame({"A": [1, 2, 3]})
psdf2 = ps.DataFrame({("A", "B"): [4, 5, 6]})
with self.assertRaisesRegex(
ValueError, "cannot join with no overlapping index names"):
np.left_shift(psdf1, psdf2)
def test_sql_with_pandas_on_spark_objects(self):
psdf = ps.DataFrame({"a": [1, 2, 3, 4]})
self.assert_eq(ps.sql("SELECT {col} FROM {tbl}", col=psdf.a, tbl=psdf),
psdf)
self.assert_eq(ps.sql("SELECT {tbl.a} FROM {tbl}", tbl=psdf), psdf)
psdf = ps.DataFrame({"A": [1, 2, 3], "B": [4, 5, 6]})
self.assert_eq(
ps.sql("SELECT {col}, {col2} FROM {tbl}",
col=psdf.A,
col2=psdf.B,
tbl=psdf), psdf)
self.assert_eq(ps.sql("SELECT {tbl.A}, {tbl.B} FROM {tbl}", tbl=psdf),
psdf)
def test_spark_io(self):
with self.temp_dir() as tmp:
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
# Write out partitioned by one column
expected.to_spark_io(tmp,
format="json",
mode="overwrite",
partition_cols="i32")
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_spark_io(tmp, format="json")
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Write out partitioned by two columns
expected.to_spark_io(tmp,
format="json",
mode="overwrite",
partition_cols=["i32", "bhello"])
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_spark_io(path=tmp, format="json")
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# When index columns are known
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
col_order = ["f", "i32", "i64"]
expected_idx = expected.set_index("bhello")[col_order]
actual_idx = ps.read_spark_io(tmp,
format="json",
index_col="bhello")[col_order]
self.assert_eq(
actual_idx.sort_values(by="f").to_spark().toPandas(),
expected_idx.sort_values(by="f").to_spark().toPandas(),
)
def test_compute_hist_multi_columns(self):
expected_bins = np.linspace(1, 50, 11)
kdf = ps.DataFrame(
{
"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50],
"b": [50, 50, 30, 30, 30, 24, 10, 5, 4, 3, 1],
}
)
bins = HistogramPlotBase.get_bins(kdf.to_spark(), 10)
self.assert_eq(pd.Series(expected_bins), pd.Series(bins))
expected_histograms = [
np.array([5, 4, 1, 0, 0, 0, 0, 0, 0, 1]),
np.array([4, 1, 0, 0, 1, 3, 0, 0, 0, 2]),
]
histograms = HistogramPlotBase.compute_hist(kdf, bins)
expected_names = ["a", "b"]
for histogram, expected_histogram, expected_name in zip(
histograms, expected_histograms, expected_names
):
self.assert_eq(
pd.Series(expected_histogram, name=expected_name), histogram, almost=True
)
def test_timestamp_subtraction(self):
pdf = self.pdf1
kdf = ps.from_pandas(pdf)
# Those fail in certain OSs presumably due to different
# timezone behaviours inherited from C library.
actual = (kdf["end_date"] - kdf["start_date"] - 1).to_pandas()
expected = (pdf["end_date"] - pdf["start_date"]) // np.timedelta64(1, "s") - 1
# self.assert_eq(actual, expected)
actual = (kdf["end_date"] - pd.Timestamp("2012-1-1 12:45:31") - 1).to_pandas()
expected = (pdf["end_date"] - pd.Timestamp("2012-1-1 12:45:31")) // np.timedelta64(
1, "s"
) - 1
# self.assert_eq(actual, expected)
actual = (pd.Timestamp("2013-3-11 21:45:00") - kdf["start_date"] - 1).to_pandas()
expected = (pd.Timestamp("2013-3-11 21:45:00") - pdf["start_date"]) // np.timedelta64(
1, "s"
) - 1
# self.assert_eq(actual, expected)
kdf = ps.DataFrame(
{"a": pd.date_range("2016-12-31", "2017-01-08", freq="D"), "b": pd.Series(range(9))}
)
expected_error_message = "datetime subtraction can only be applied to datetime series."
with self.assertRaisesRegex(TypeError, expected_error_message):
kdf["a"] - kdf["b"]
with self.assertRaisesRegex(TypeError, expected_error_message):
kdf["a"] - 1
with self.assertRaisesRegex(TypeError, expected_error_message):
1 - kdf["a"]
def test_to_csv_with_partition_cols(self):
pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]})
kdf = ps.DataFrame(pdf)
kdf.to_csv(self.tmp_dir, partition_cols="b", num_files=1)
partition_paths = [
path for path in os.listdir(self.tmp_dir) if path.startswith("b=")
]
assert len(partition_paths) > 0
for partition_path in partition_paths:
column, value = partition_path.split("=")
expected = pdf[pdf[column] == value].drop(
"b", axis=1).to_csv(index=False)
output_paths = [
path for path in os.listdir("%s/%s" %
(self.tmp_dir, partition_path))
if path.startswith("part-")
]
assert len(output_paths) > 0
output_path = "%s/%s/%s" % (self.tmp_dir, partition_path,
output_paths[0])
with open(output_path) as f:
self.assertEqual(f.read(), expected)
def test_parquet_read(self):
with self.temp_dir() as tmp:
data = self.test_pdf
self.spark.createDataFrame(
data, "i32 int, i64 long, f double, bhello string").coalesce(
1).write.parquet(tmp, mode="overwrite")
def check(columns):
expected = pd.read_parquet(tmp, columns=columns)
actual = ps.read_parquet(tmp, columns=columns)
self.assertPandasEqual(expected, actual.to_pandas())
check(None)
check(["i32", "i64"])
check(["i64", "i32"])
# check with pyspark patch.
expected = pd.read_parquet(tmp)
actual = ps.read_parquet(tmp)
self.assertPandasEqual(expected, actual.to_pandas())
# When index columns are known
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]]
actual_idx = ps.read_parquet(
tmp, index_col="bhello")[["f", "i32", "i64"]]
self.assert_eq(
actual_idx.sort_values(by="f").to_spark().toPandas(),
expected_idx.sort_values(by="f").to_spark().toPandas(),
)
def _convert_var(self, var: Any) -> Any:
"""
Converts a python object into a string that is legal SQL.
"""
if isinstance(var, (int, float)):
return str(var)
if isinstance(var, Series):
return self._convert_var(var.to_dataframe())
if isinstance(var, pd.DataFrame):
return self._convert_var(ps.DataFrame(var))
if isinstance(var, DataFrame):
df_id = "pandas_on_spark_" + str(id(var))
if df_id not in self._temp_views:
sdf = var._to_spark()
sdf.createOrReplaceTempView(df_id)
self._temp_views[df_id] = sdf
return df_id
if isinstance(var, str):
return '"' + escape_sql_string(var) + '"'
if isinstance(var, list):
return "(" + ", ".join([self._convert_var(v) for v in var]) + ")"
if isinstance(var, (tuple, range)):
return self._convert_var(list(var))
raise ValueError("Unsupported variable type {}: {}".format(
type(var).__name__, str(var)))
def test_checkpoint(self):
with self.temp_dir() as tmp:
self.spark.sparkContext.setCheckpointDir(tmp)
psdf = ps.DataFrame({"a": ["a", "b", "c"]})
new_psdf = psdf.spark.checkpoint()
self.assertIsNotNone(os.listdir(tmp))
self.assert_eq(psdf, new_psdf)
def test_orc_write(self):
with self.temp_dir() as tmp:
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
# Write out partitioned by one column
expected.to_orc(tmp, mode="overwrite", partition_cols="i32")
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_orc(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Write out partitioned by two columns
expected.to_orc(tmp,
mode="overwrite",
partition_cols=["i32", "bhello"])
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_orc(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
def test_get_index_map(self):
psdf = ps.DataFrame({
"year": [2015, 2016],
"month": [2, 3],
"day": [4, 5]
})
sdf = psdf.to_spark()
self.assertEqual(_get_index_map(sdf), (None, None))
def check(actual, expected):
actual_scols, actual_labels = actual
expected_column_names, expected_labels = expected
self.assertEqual(len(actual_scols), len(expected_column_names))
for actual_scol, expected_column_name in zip(
actual_scols, expected_column_names):
expected_scol = sdf[expected_column_name]
self.assertTrue(spark_column_equals(actual_scol,
expected_scol))
self.assertEqual(actual_labels, expected_labels)
check(_get_index_map(sdf, "year"), (["year"], [("year", )]))
check(_get_index_map(sdf, ["year", "month"]),
(["year", "month"], [("year", ), ("month", )]))
self.assertRaises(KeyError,
lambda: _get_index_map(sdf, ["year", "hour"]))
def test_date_subtraction(self):
pdf = self.pdf1
kdf = ps.from_pandas(pdf)
self.assert_eq(
kdf["end_date"].dt.date - kdf["start_date"].dt.date,
(pdf["end_date"].dt.date - pdf["start_date"].dt.date).dt.days,
)
self.assert_eq(
kdf["end_date"].dt.date - datetime.date(2012, 1, 1),
(pdf["end_date"].dt.date - datetime.date(2012, 1, 1)).dt.days,
)
self.assert_eq(
datetime.date(2013, 3, 11) - kdf["start_date"].dt.date,
(datetime.date(2013, 3, 11) - pdf["start_date"].dt.date).dt.days,
)
kdf = ps.DataFrame({
"a":
pd.date_range("2016-12-31", "2017-01-08", freq="D"),
"b":
pd.Series(range(9))
})
expected_error_message = "date subtraction can only be applied to date series."
with self.assertRaisesRegex(TypeError, expected_error_message):
kdf["a"].dt.date - kdf["b"]
with self.assertRaisesRegex(TypeError, expected_error_message):
kdf["a"].dt.date - 1
with self.assertRaisesRegex(TypeError, expected_error_message):
1 - kdf["a"].dt.date
def test_kde_plot(self):
psdf = ps.DataFrame({
"a": [1, 2, 3, 4, 5],
"b": [1, 3, 5, 7, 9],
"c": [2, 4, 6, 8, 10]
})
pdf = pd.DataFrame({
"Density": [
0.03515491,
0.06834979,
0.00663503,
0.02372059,
0.06834979,
0.01806934,
0.01806934,
0.06834979,
0.02372059,
],
"names": ["a", "a", "a", "b", "b", "b", "c", "c", "c"],
"index": [-3.5, 5.5, 14.5, -3.5, 5.5, 14.5, -3.5, 5.5, 14.5],
})
actual = psdf.plot.kde(bw_method=5, ind=3)
expected = express.line(pdf, x="index", y="Density", color="names")
expected["layout"]["xaxis"]["title"] = None
self.assertEqual(pprint.pformat(actual.to_dict()),
pprint.pformat(expected.to_dict()))
def test_table(self):
with self.table("test_table"):
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
# Write out partitioned by one column
expected.spark.to_table("test_table",
mode="overwrite",
partition_cols="i32")
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_table("test_table")
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Write out partitioned by two columns
expected.to_table("test_table",
mode="overwrite",
partition_cols=["i32", "bhello"])
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_table("test_table")
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# When index columns are known
expected_idx = expected.set_index("bhello")[["f", "i32", "i64"]]
actual_idx = ps.read_table("test_table",
index_col="bhello")[["f", "i32", "i64"]]
self.assert_eq(
actual_idx.sort_values(by="f").to_spark().toPandas(),
expected_idx.sort_values(by="f").to_spark().toPandas(),
)
expected_idx = expected.set_index(["bhello"])[["f", "i32", "i64"]]
actual_idx = ps.read_table("test_table",
index_col=["bhello"
])[["f", "i32", "i64"]]
self.assert_eq(
actual_idx.sort_values(by="f").to_spark().toPandas(),
expected_idx.sort_values(by="f").to_spark().toPandas(),
)
expected_idx = expected.set_index(["i32", "bhello"])[["f", "i64"]]
actual_idx = ps.read_table("test_table",
index_col=["i32",
"bhello"])[["f", "i64"]]
self.assert_eq(
actual_idx.sort_values(by="f").to_spark().toPandas(),
expected_idx.sort_values(by="f").to_spark().toPandas(),
)
def test_to_csv_with_path_and_basic_options(self):
pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": ["a", "b", "c"]})
kdf = ps.DataFrame(pdf)
kdf.to_csv(self.tmp_dir, num_files=1, sep="|", header=False, columns=["aa"])
expected = pdf.to_csv(index=False, sep="|", header=False, columns=["aa"])
self._check_output(self.tmp_dir, expected)
def test_to_csv_with_path_and_pyspark_options(self):
pdf = pd.DataFrame({"a": [1, 2, 3, None], "b": ["a", "b", "c", None]})
kdf = ps.DataFrame(pdf)
kdf.to_csv(self.tmp_dir, nullValue="null", num_files=1)
expected = pdf.to_csv(index=False, na_rep="null")
self._check_output(self.tmp_dir, expected)
def test_missing(self):
kdf = ps.DataFrame(np.random.rand(2500, 4), columns=["a", "b", "c", "d"])
unsupported_functions = ["box", "hexbin"]
for name in unsupported_functions:
with self.assertRaisesRegex(
PandasNotImplementedError, "method.*DataFrame.*{}.*not implemented".format(name)
):
getattr(kdf.plot, name)()
def test_to_json_with_path(self):
pdf = pd.DataFrame({"a": [1], "b": ["a"]})
psdf = ps.DataFrame(pdf)
psdf.to_json(self.tmp_dir, num_files=1)
expected = pdf.to_json(orient="records")
output_paths = [path for path in os.listdir(self.tmp_dir) if path.startswith("part-")]
assert len(output_paths) > 0
output_path = "%s/%s" % (self.tmp_dir, output_paths[0])
self.assertEqual("[%s]" % open(output_path).read().strip(), expected)
def test_compute_hist_single_column(self):
kdf = ps.DataFrame(
{"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 15, 50]}, index=[0, 1, 3, 5, 6, 8, 9, 9, 9, 10, 10]
)
expected_bins = np.linspace(1, 50, 11)
bins = HistogramPlotBase.get_bins(kdf[["a"]].to_spark(), 10)
expected_histogram = np.array([5, 4, 1, 0, 0, 0, 0, 0, 0, 1])
histogram = HistogramPlotBase.compute_hist(kdf[["a"]], bins)[0]
self.assert_eq(pd.Series(expected_bins), pd.Series(bins))
self.assert_eq(pd.Series(expected_histogram, name="a"), histogram, almost=True)
def test_to_csv_with_path_and_basic_options_multiindex_columns(self):
pdf = pd.DataFrame({("x", "a"): [1, 2, 3], ("y", "b"): ["a", "b", "c"]})
kdf = ps.DataFrame(pdf)
with self.assertRaises(ValueError):
kdf.to_csv(self.tmp_dir, num_files=1, sep="|", columns=[("x", "a")])
kdf.to_csv(self.tmp_dir, num_files=1, sep="|", header=["a"], columns=[("x", "a")])
pdf.columns = ["a", "b"]
expected = pdf.to_csv(index=False, sep="|", columns=["a"])
self._check_output(self.tmp_dir, expected)
def test_repartition(self):
kdf = ps.DataFrame({
"age": [5, 5, 2, 2],
"name": ["Bob", "Bob", "Alice", "Alice"]
})
num_partitions = kdf.to_spark().rdd.getNumPartitions() + 1
num_partitions += 1
new_kdf = kdf.spark.repartition(num_partitions)
self.assertEqual(new_kdf.to_spark().rdd.getNumPartitions(),
num_partitions)
self.assert_eq(kdf.sort_index(), new_kdf.sort_index())
# Reserves Index
kdf = kdf.set_index("age")
num_partitions += 1
new_kdf = kdf.spark.repartition(num_partitions)
self.assertEqual(new_kdf.to_spark().rdd.getNumPartitions(),
num_partitions)
self.assert_eq(kdf.sort_index(), new_kdf.sort_index())
# Reflects internal changes
kdf = kdf.reset_index()
kdf = kdf.set_index("name")
kdf2 = kdf + 1
num_partitions += 1
self.assert_eq(kdf2.sort_index(),
(kdf +
1).spark.repartition(num_partitions).sort_index())
# Reserves MultiIndex
kdf = ps.DataFrame({"a": ["a", "b", "c"]},
index=[[1, 2, 3], [4, 5, 6]])
num_partitions = kdf.to_spark().rdd.getNumPartitions() + 1
new_kdf = kdf.spark.repartition(num_partitions)
self.assertEqual(new_kdf.to_spark().rdd.getNumPartitions(),
num_partitions)
self.assert_eq(kdf.sort_index(), new_kdf.sort_index())
def test_repr_float_index(self):
kdf = ps.DataFrame(
{"a": np.random.rand(ReprTest.max_display_count)},
index=np.random.rand(ReprTest.max_display_count),
)
self.assertTrue("Showing only the first" not in repr(kdf))
self.assert_eq(repr(kdf), repr(kdf.to_pandas()))
self.assertTrue("Showing only the first" not in repr(kdf.a))
self.assert_eq(repr(kdf.a), repr(kdf.a.to_pandas()))
self.assertTrue("Showing only the first" not in repr(kdf.index))
self.assert_eq(repr(kdf.index), repr(kdf.index.to_pandas()))
self.assertTrue("Showing only the first" not in kdf._repr_html_())
self.assertEqual(kdf._repr_html_(), kdf.to_pandas()._repr_html_())
kdf = ps.DataFrame(
{"a": np.random.rand(ReprTest.max_display_count + 1)},
index=np.random.rand(ReprTest.max_display_count + 1),
)
self.assertTrue("Showing only the first" in repr(kdf))
self.assertTrue("Showing only the first" in repr(kdf.a))
self.assertTrue("Showing only the first" in repr(kdf.index))
self.assertTrue("Showing only the first" in kdf._repr_html_())
def test_parquet_write(self):
with self.temp_dir() as tmp:
pdf = self.test_pdf
expected = ps.DataFrame(pdf)
# Write out partitioned by one column
expected.to_parquet(tmp, mode="overwrite", partition_cols="i32")
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_parquet(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Write out partitioned by two columns
expected.to_parquet(tmp, mode="overwrite", partition_cols=["i32", "bhello"])
# Reset column order, as once the data is written out, Spark rearranges partition
# columns to appear first.
actual = ps.read_parquet(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Set `compression` with string
expected.to_parquet(tmp, mode="overwrite", partition_cols="i32", compression="none")
actual = ps.read_parquet(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
# Test `options` parameter
expected.to_parquet(
tmp, mode="overwrite", partition_cols="i32", options={"compression": "none"}
)
actual = ps.read_parquet(tmp)
self.assertFalse((actual.columns == self.test_column_order).all())
actual = actual[self.test_column_order]
self.assert_eq(
actual.sort_values(by="f").to_spark().toPandas(),
expected.sort_values(by="f").to_spark().toPandas(),
)
def test_timestamp_subtraction_errors(self):
psdf = ps.DataFrame({
"a":
pd.date_range("2016-12-31", "2017-01-08", freq="D"),
"b":
pd.Series(range(9))
})
expected_error_message = "Datetime subtraction can only be applied to datetime series."
with self.assertRaisesRegex(TypeError, expected_error_message):
psdf["a"] - psdf["b"]
with self.assertRaisesRegex(TypeError, expected_error_message):
psdf["a"] - 1
with self.assertRaisesRegex(TypeError, expected_error_message):
1 - psdf["a"]