| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| import os |
| from contextlib import contextmanager |
| |
| import pandas as pd |
| import numpy as np |
| |
| from pyspark import pandas as ps |
| from pyspark.testing.pandasutils import PandasOnSparkTestCase, TestUtils |
| |
| |
| def normalize_text(s): |
| return "\n".join(map(str.strip, s.strip().split("\n"))) |
| |
| |
| class CsvTestsMixin: |
| @property |
| def csv_text(self): |
| return normalize_text( |
| """ |
| name,amount |
| Alice,100 |
| Bob,-200 |
| Charlie,300 |
| Dennis,400 |
| Edith,-500 |
| Frank,600 |
| Alice,200 |
| Frank,-200 |
| Bob,600 |
| Alice,400 |
| Frank,200 |
| Alice,300 |
| Edith,600 |
| """ |
| ) |
| |
| @property |
| def csv_text_2(self): |
| return normalize_text( |
| """ |
| A,B |
| item1,1 |
| item2,1,2 |
| item3,1,2,3,4 |
| item4,1 |
| """ |
| ) |
| |
| @property |
| def csv_text_with_comments(self): |
| return normalize_text( |
| """ |
| # header |
| %s |
| # comment |
| Alice,400 |
| Edith,600 |
| # footer |
| """ |
| % self.csv_text |
| ) |
| |
| @property |
| def tab_delimited_csv_text(self): |
| return normalize_text( |
| """ |
| name\tamount |
| Alice\t100 |
| Bob\t-200 |
| Charlie\t300 |
| """ |
| ) |
| |
| @property |
| def q_quoted_csv_text(self): |
| return normalize_text( |
| """ |
| QnameQ,QamountQ |
| QA,liceQ,Q100Q |
| QB,obQ,Q-200Q |
| QC,harlieQ,Q300Q |
| """ |
| ) |
| |
| @property |
| def e_escapeted_csv_text(self): |
| return normalize_text( |
| """ |
| name,amount |
| "AE"lice",100 |
| "BE"ob",-200 |
| "CE"harlie",300 |
| """ |
| ) |
| |
| @contextmanager |
| def csv_file(self, csv): |
| with self.temp_file() as tmp: |
| with open(tmp, "w") as f: |
| f.write(csv) |
| yield tmp |
| |
| def test_read_csv(self): |
| with self.csv_file(self.csv_text) as fn: |
| |
| def check(header="infer", names=None, usecols=None, index_col=None): |
| expected = pd.read_csv( |
| fn, header=header, names=names, usecols=usecols, index_col=index_col |
| ) |
| actual = ps.read_csv( |
| fn, header=header, names=names, usecols=usecols, index_col=index_col |
| ) |
| self.assert_eq(expected, actual, almost=True) |
| |
| check() |
| check(header=0) |
| check(header=None) |
| check(names=["n", "a"]) |
| check(names=[("x", "n"), ("y", "a")]) |
| check(names=[10, 20]) |
| check(header=0, names=["n", "a"]) |
| check(usecols=[1]) |
| check(usecols=[1, 0]) |
| check(usecols=["amount"]) |
| check(usecols=["amount", "name"]) |
| check(usecols=[]) |
| check(usecols=[1, 1]) |
| check(usecols=["amount", "amount"]) |
| check(header=None, usecols=[1]) |
| check(names=["n", "a"], usecols=["a"]) |
| check(header=None, names=["n", "a"], usecols=["a"]) |
| check(index_col=["amount"]) |
| check(header=None, index_col=[1]) |
| check(names=["n", "a"], index_col=["a"]) |
| check(names=["n", "a"], index_col="a") |
| |
| # check with pyspark patch. |
| expected = pd.read_csv(fn) |
| actual = ps.read_csv(fn) |
| self.assert_eq(expected, actual, almost=True) |
| |
| self.assertRaisesRegex( |
| ValueError, "non-unique", lambda: ps.read_csv(fn, names=["n", "n"]) |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "does not match the number.*3", |
| lambda: ps.read_csv(fn, names=["n", "a", "b"]), |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "does not match the number.*3", |
| lambda: ps.read_csv(fn, header=0, names=["n", "a", "b"]), |
| ) |
| self.assertRaisesRegex( |
| ValueError, "Usecols do not match.*3", lambda: ps.read_csv(fn, usecols=[1, 3]) |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "Usecols do not match.*col", |
| lambda: ps.read_csv(fn, usecols=["amount", "col"]), |
| ) |
| self.assertRaisesRegex( |
| ValueError, "Unknown header argument 1", lambda: ps.read_csv(fn, header="1") |
| ) |
| expected_error_message = ( |
| "'usecols' must either be list-like of all strings, " |
| "all unicode, all integers or a callable." |
| ) |
| self.assertRaisesRegex( |
| ValueError, expected_error_message, lambda: ps.read_csv(fn, usecols=[1, "amount"]) |
| ) |
| |
| # check with index_col |
| expected = pd.read_csv(fn).set_index("name") |
| actual = ps.read_csv(fn, index_col="name") |
| self.assert_eq(expected, actual, almost=True) |
| |
| def test_read_with_spark_schema(self): |
| with self.csv_file(self.csv_text_2) as fn: |
| actual = ps.read_csv(fn, names="A string, B string, C long, D long, E long") |
| expected = pd.read_csv(fn, names=["A", "B", "C", "D", "E"]) |
| self.assert_eq(expected, actual) |
| |
| def test_read_csv_with_comment(self): |
| with self.csv_file(self.csv_text_with_comments) as fn: |
| expected = pd.read_csv(fn, comment="#") |
| actual = ps.read_csv(fn, comment="#") |
| self.assert_eq(expected, actual, almost=True) |
| |
| self.assertRaisesRegex( |
| ValueError, |
| "Only length-1 comment characters supported", |
| lambda: ps.read_csv(fn, comment="").show(), |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "Only length-1 comment characters supported", |
| lambda: ps.read_csv(fn, comment="##").show(), |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "Only length-1 comment characters supported", |
| lambda: ps.read_csv(fn, comment=1), |
| ) |
| self.assertRaisesRegex( |
| ValueError, |
| "Only length-1 comment characters supported", |
| lambda: ps.read_csv(fn, comment=[1]), |
| ) |
| |
| def test_read_csv_with_limit(self): |
| with self.csv_file(self.csv_text_with_comments) as fn: |
| expected = pd.read_csv(fn, comment="#", nrows=2) |
| actual = ps.read_csv(fn, comment="#", nrows=2) |
| self.assert_eq(expected, actual, almost=True) |
| |
| def test_read_csv_with_encoding(self): |
| # SPARK-37181: Read csv supporting latin-1 encoding. |
| with self.csv_file(self.csv_text) as fn: |
| expected = pd.read_csv(fn, encoding="latin-1") |
| actual = ps.read_csv(fn, encoding="latin-1") |
| self.assert_eq(expected, actual, almost=True) |
| |
| def test_read_csv_with_sep(self): |
| with self.csv_file(self.tab_delimited_csv_text) as fn: |
| expected = pd.read_csv(fn, sep="\t") |
| actual = ps.read_csv(fn, sep="\t") |
| self.assert_eq(expected, actual, almost=True) |
| |
| def test_read_csv_with_parse_dates(self): |
| self.assertRaisesRegex( |
| ValueError, "parse_dates", lambda: ps.read_csv("path", parse_dates=True) |
| ) |
| |
| def test_read_csv_with_dtype(self): |
| with self.csv_file(self.csv_text) as fn: |
| self.assert_eq(ps.read_csv(fn), pd.read_csv(fn), almost=True) |
| self.assert_eq(ps.read_csv(fn, dtype=str), pd.read_csv(fn, dtype=str)) |
| self.assert_eq( |
| ps.read_csv(fn, dtype={"amount": "int64"}), |
| pd.read_csv(fn, dtype={"amount": "int64"}), |
| ) |
| |
| def test_read_csv_with_quotechar(self): |
| with self.csv_file(self.q_quoted_csv_text) as fn: |
| self.assert_eq( |
| ps.read_csv(fn, quotechar="Q"), pd.read_csv(fn, quotechar="Q"), almost=True |
| ) |
| |
| def test_read_csv_with_escapechar(self): |
| with self.csv_file(self.e_escapeted_csv_text) as fn: |
| self.assert_eq( |
| ps.read_csv(fn, escapechar="E"), pd.read_csv(fn, escapechar="E"), almost=True |
| ) |
| |
| self.assert_eq( |
| ps.read_csv(fn, escapechar="ABC", escape="E"), |
| pd.read_csv(fn, escapechar="E"), |
| almost=True, |
| ) |
| |
| def test_to_csv(self): |
| pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": [4, 5, 6]}, index=[0, 1, 3]) |
| psdf = ps.DataFrame(pdf) |
| |
| self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) |
| self.assert_eq(psdf.to_csv(columns=["aa"]), pdf.to_csv(columns=["aa"], index=False)) |
| self.assert_eq(psdf.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]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.to_csv(na_rep="null"), pdf.to_csv(na_rep="null", index=False)) |
| self.assert_eq( |
| psdf.a.to_csv(na_rep="null"), pdf.a.to_csv(na_rep="null", index=False, header=True) |
| ) |
| |
| self.assertRaises(KeyError, lambda: psdf.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]) |
| psdf = ps.from_pandas(pdf) |
| |
| self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) |
| self.assert_eq(psdf.to_csv(header=False), pdf.to_csv(header=False, index=False)) |
| self.assert_eq(psdf.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]) |
| psdf = ps.DataFrame(pdf) |
| |
| self.assert_eq(psdf.to_csv(), pdf.to_csv(index=False)) |
| self.assert_eq(psdf.to_csv(columns=[10]), pdf.to_csv(columns=[10], index=False)) |
| |
| self.assertRaises(TypeError, lambda: psdf.to_csv(columns=10)) |
| |
| def _check_output(self, dir, expected): |
| output_paths = [path for path in os.listdir(dir) if path.startswith("part-")] |
| assert len(output_paths) > 0 |
| output_path = "%s/%s" % (dir, output_paths[0]) |
| with open(output_path) as f: |
| self.assertEqual(f.read(), expected) |
| |
| def test_to_csv_with_path(self): |
| pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| psdf = ps.DataFrame(pdf) |
| |
| with self.temp_dir() as dirpath: |
| tmp_dir = "{}/tmp1".format(dirpath) |
| |
| psdf.to_csv(tmp_dir, num_files=1) |
| self._check_output(tmp_dir, pdf.to_csv(index=False)) |
| |
| tmp_dir = "{}/tmp2".format(dirpath) |
| |
| self.assertRaises(KeyError, lambda: psdf.to_csv(tmp_dir, columns=["c"], num_files=1)) |
| |
| # non-string names |
| pdf = pd.DataFrame({10: [1, 2, 3], 20: ["a", "b", "c"]}) |
| psdf = ps.DataFrame(pdf) |
| |
| tmp_dir = "{}/tmp3".format(dirpath) |
| |
| psdf.to_csv(tmp_dir, num_files=1) |
| self._check_output(tmp_dir, pdf.to_csv(index=False)) |
| |
| tmp_dir = "{}/tmp4".format(dirpath) |
| |
| psdf.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(dirpath) |
| |
| self.assertRaises(TypeError, lambda: psdf.to_csv(tmp_dir, columns=10, num_files=1)) |
| |
| def test_to_csv_with_path_and_basic_options(self): |
| pdf = pd.DataFrame({"aa": [1, 2, 3], "bb": ["a", "b", "c"]}) |
| psdf = ps.DataFrame(pdf) |
| |
| with self.temp_dir() as dirpath: |
| psdf.to_csv(dirpath, num_files=1, sep="|", header=False, columns=["aa"]) |
| expected = pdf.to_csv(index=False, sep="|", header=False, columns=["aa"]) |
| |
| self._check_output(dirpath, expected) |
| |
| 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"]}) |
| psdf = ps.DataFrame(pdf) |
| |
| with self.temp_dir() as dirpath: |
| with self.assertRaises(ValueError): |
| psdf.to_csv(dirpath, num_files=1, sep="|", columns=[("x", "a")]) |
| |
| psdf.to_csv(dirpath, 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(dirpath, expected) |
| |
| def test_to_csv_with_path_and_pyspark_options(self): |
| pdf = pd.DataFrame({"a": [1, 2, 3, None], "b": ["a", "b", "c", None]}) |
| psdf = ps.DataFrame(pdf) |
| |
| with self.temp_dir() as dirpath: |
| psdf.to_csv(dirpath, nullValue="null", num_files=1) |
| expected = pdf.to_csv(index=False, na_rep="null") |
| |
| self._check_output(dirpath, expected) |
| |
| def test_to_csv_with_partition_cols(self): |
| pdf = pd.DataFrame({"a": [1, 2, 3], "b": ["a", "b", "c"]}) |
| psdf = ps.DataFrame(pdf) |
| |
| with self.temp_dir() as dirpath: |
| psdf.to_csv(dirpath, partition_cols="b", num_files=1) |
| |
| partition_paths = [path for path in os.listdir(dirpath) 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" % (dirpath, partition_path)) |
| if path.startswith("part-") |
| ] |
| assert len(output_paths) > 0 |
| output_path = "%s/%s/%s" % (dirpath, partition_path, output_paths[0]) |
| with open(output_path) as f: |
| self.assertEqual(f.read(), expected) |
| |
| |
| class CsvTests( |
| CsvTestsMixin, |
| PandasOnSparkTestCase, |
| TestUtils, |
| ): |
| pass |
| |
| |
| if __name__ == "__main__": |
| import unittest |
| from pyspark.pandas.tests.io.test_csv import * # noqa: F401 |
| |
| try: |
| import xmlrunner |
| |
| testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2) |
| except ImportError: |
| testRunner = None |
| unittest.main(testRunner=testRunner, verbosity=2) |
