--- /dev/null
+# 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
+import pickle
+import random
+import unittest
+from io import BytesIO
+from os.path import join as pjoin
+
+import numpy as np
+import pytest
+
+import pyarrow as pa
+from pyarrow.pandas_compat import _pandas_api
+from pyarrow.tests import util
+from pyarrow.tests.parquet.common import _test_dataframe
+from pyarrow.tests.parquet.test_dataset import (
+ _test_read_common_metadata_files, _test_write_to_dataset_with_partitions,
+ _test_write_to_dataset_no_partitions
+)
+from pyarrow.util import guid
+
+# ----------------------------------------------------------------------
+# HDFS tests
+
+
+def check_libhdfs_present():
+ if not pa.have_libhdfs():
+ message = 'No libhdfs available on system'
+ if os.environ.get('PYARROW_HDFS_TEST_LIBHDFS_REQUIRE'):
+ pytest.fail(message)
+ else:
+ pytest.skip(message)
+
+
+def hdfs_test_client():
+ host = os.environ.get('ARROW_HDFS_TEST_HOST', 'default')
+ user = os.environ.get('ARROW_HDFS_TEST_USER', None)
+ try:
+ port = int(os.environ.get('ARROW_HDFS_TEST_PORT', 0))
+ except ValueError:
+ raise ValueError('Env variable ARROW_HDFS_TEST_PORT was not '
+ 'an integer')
+
+ with pytest.warns(FutureWarning):
+ return pa.hdfs.connect(host, port, user)
+
+
+@pytest.mark.hdfs
+class HdfsTestCases:
+
+ def _make_test_file(self, hdfs, test_name, test_path, test_data):
+ base_path = pjoin(self.tmp_path, test_name)
+ hdfs.mkdir(base_path)
+
+ full_path = pjoin(base_path, test_path)
+
+ with hdfs.open(full_path, 'wb') as f:
+ f.write(test_data)
+
+ return full_path
+
+ @classmethod
+ def setUpClass(cls):
+ cls.check_driver()
+ cls.hdfs = hdfs_test_client()
+ cls.tmp_path = '/tmp/pyarrow-test-{}'.format(random.randint(0, 1000))
+ cls.hdfs.mkdir(cls.tmp_path)
+
+ @classmethod
+ def tearDownClass(cls):
+ cls.hdfs.delete(cls.tmp_path, recursive=True)
+ cls.hdfs.close()
+
+ def test_pickle(self):
+ s = pickle.dumps(self.hdfs)
+ h2 = pickle.loads(s)
+ assert h2.is_open
+ assert h2.host == self.hdfs.host
+ assert h2.port == self.hdfs.port
+ assert h2.user == self.hdfs.user
+ assert h2.kerb_ticket == self.hdfs.kerb_ticket
+ # smoketest unpickled client works
+ h2.ls(self.tmp_path)
+
+ def test_cat(self):
+ path = pjoin(self.tmp_path, 'cat-test')
+
+ data = b'foobarbaz'
+ with self.hdfs.open(path, 'wb') as f:
+ f.write(data)
+
+ contents = self.hdfs.cat(path)
+ assert contents == data
+
+ def test_capacity_space(self):
+ capacity = self.hdfs.get_capacity()
+ space_used = self.hdfs.get_space_used()
+ disk_free = self.hdfs.df()
+
+ assert capacity > 0
+ assert capacity > space_used
+ assert disk_free == (capacity - space_used)
+
+ def test_close(self):
+ client = hdfs_test_client()
+ assert client.is_open
+ client.close()
+ assert not client.is_open
+
+ with pytest.raises(Exception):
+ client.ls('/')
+
+ def test_mkdir(self):
+ path = pjoin(self.tmp_path, 'test-dir/test-dir')
+ parent_path = pjoin(self.tmp_path, 'test-dir')
+
+ self.hdfs.mkdir(path)
+ assert self.hdfs.exists(path)
+
+ self.hdfs.delete(parent_path, recursive=True)
+ assert not self.hdfs.exists(path)
+
+ def test_mv_rename(self):
+ path = pjoin(self.tmp_path, 'mv-test')
+ new_path = pjoin(self.tmp_path, 'mv-new-test')
+
+ data = b'foobarbaz'
+ with self.hdfs.open(path, 'wb') as f:
+ f.write(data)
+
+ assert self.hdfs.exists(path)
+ self.hdfs.mv(path, new_path)
+ assert not self.hdfs.exists(path)
+ assert self.hdfs.exists(new_path)
+
+ assert self.hdfs.cat(new_path) == data
+
+ self.hdfs.rename(new_path, path)
+ assert self.hdfs.cat(path) == data
+
+ def test_info(self):
+ path = pjoin(self.tmp_path, 'info-base')
+ file_path = pjoin(path, 'ex')
+ self.hdfs.mkdir(path)
+
+ data = b'foobarbaz'
+ with self.hdfs.open(file_path, 'wb') as f:
+ f.write(data)
+
+ path_info = self.hdfs.info(path)
+ file_path_info = self.hdfs.info(file_path)
+
+ assert path_info['kind'] == 'directory'
+
+ assert file_path_info['kind'] == 'file'
+ assert file_path_info['size'] == len(data)
+
+ def test_exists_isdir_isfile(self):
+ dir_path = pjoin(self.tmp_path, 'info-base')
+ file_path = pjoin(dir_path, 'ex')
+ missing_path = pjoin(dir_path, 'this-path-is-missing')
+
+ self.hdfs.mkdir(dir_path)
+ with self.hdfs.open(file_path, 'wb') as f:
+ f.write(b'foobarbaz')
+
+ assert self.hdfs.exists(dir_path)
+ assert self.hdfs.exists(file_path)
+ assert not self.hdfs.exists(missing_path)
+
+ assert self.hdfs.isdir(dir_path)
+ assert not self.hdfs.isdir(file_path)
+ assert not self.hdfs.isdir(missing_path)
+
+ assert not self.hdfs.isfile(dir_path)
+ assert self.hdfs.isfile(file_path)
+ assert not self.hdfs.isfile(missing_path)
+
+ def test_disk_usage(self):
+ path = pjoin(self.tmp_path, 'disk-usage-base')
+ p1 = pjoin(path, 'p1')
+ p2 = pjoin(path, 'p2')
+
+ subdir = pjoin(path, 'subdir')
+ p3 = pjoin(subdir, 'p3')
+
+ if self.hdfs.exists(path):
+ self.hdfs.delete(path, True)
+
+ self.hdfs.mkdir(path)
+ self.hdfs.mkdir(subdir)
+
+ data = b'foobarbaz'
+
+ for file_path in [p1, p2, p3]:
+ with self.hdfs.open(file_path, 'wb') as f:
+ f.write(data)
+
+ assert self.hdfs.disk_usage(path) == len(data) * 3
+
+ def test_ls(self):
+ base_path = pjoin(self.tmp_path, 'ls-test')
+ self.hdfs.mkdir(base_path)
+
+ dir_path = pjoin(base_path, 'a-dir')
+ f1_path = pjoin(base_path, 'a-file-1')
+
+ self.hdfs.mkdir(dir_path)
+
+ f = self.hdfs.open(f1_path, 'wb')
+ f.write(b'a' * 10)
+
+ contents = sorted(self.hdfs.ls(base_path, False))
+ assert contents == [dir_path, f1_path]
+
+ def test_chmod_chown(self):
+ path = pjoin(self.tmp_path, 'chmod-test')
+ with self.hdfs.open(path, 'wb') as f:
+ f.write(b'a' * 10)
+
+ def test_download_upload(self):
+ base_path = pjoin(self.tmp_path, 'upload-test')
+
+ data = b'foobarbaz'
+ buf = BytesIO(data)
+ buf.seek(0)
+
+ self.hdfs.upload(base_path, buf)
+
+ out_buf = BytesIO()
+ self.hdfs.download(base_path, out_buf)
+ out_buf.seek(0)
+ assert out_buf.getvalue() == data
+
+ def test_file_context_manager(self):
+ path = pjoin(self.tmp_path, 'ctx-manager')
+
+ data = b'foo'
+ with self.hdfs.open(path, 'wb') as f:
+ f.write(data)
+
+ with self.hdfs.open(path, 'rb') as f:
+ assert f.size() == 3
+ result = f.read(10)
+ assert result == data
+
+ def test_open_not_exist(self):
+ path = pjoin(self.tmp_path, 'does-not-exist-123')
+
+ with pytest.raises(FileNotFoundError):
+ self.hdfs.open(path)
+
+ def test_open_write_error(self):
+ with pytest.raises((FileExistsError, IsADirectoryError)):
+ self.hdfs.open('/', 'wb')
+
+ def test_read_whole_file(self):
+ path = pjoin(self.tmp_path, 'read-whole-file')
+
+ data = b'foo' * 1000
+ with self.hdfs.open(path, 'wb') as f:
+ f.write(data)
+
+ with self.hdfs.open(path, 'rb') as f:
+ result = f.read()
+
+ assert result == data
+
+ def _write_multiple_hdfs_pq_files(self, tmpdir):
+ import pyarrow.parquet as pq
+ nfiles = 10
+ size = 5
+ test_data = []
+ for i in range(nfiles):
+ df = _test_dataframe(size, seed=i)
+
+ df['index'] = np.arange(i * size, (i + 1) * size)
+
+ # Hack so that we don't have a dtype cast in v1 files
+ df['uint32'] = df['uint32'].astype(np.int64)
+
+ path = pjoin(tmpdir, '{}.parquet'.format(i))
+
+ table = pa.Table.from_pandas(df, preserve_index=False)
+ with self.hdfs.open(path, 'wb') as f:
+ pq.write_table(table, f)
+
+ test_data.append(table)
+
+ expected = pa.concat_tables(test_data)
+ return expected
+
+ @pytest.mark.pandas
+ @pytest.mark.parquet
+ def test_read_multiple_parquet_files(self):
+
+ tmpdir = pjoin(self.tmp_path, 'multi-parquet-' + guid())
+
+ self.hdfs.mkdir(tmpdir)
+
+ expected = self._write_multiple_hdfs_pq_files(tmpdir)
+ result = self.hdfs.read_parquet(tmpdir)
+
+ _pandas_api.assert_frame_equal(result.to_pandas()
+ .sort_values(by='index')
+ .reset_index(drop=True),
+ expected.to_pandas())
+
+ @pytest.mark.pandas
+ @pytest.mark.parquet
+ def test_read_multiple_parquet_files_with_uri(self):
+ import pyarrow.parquet as pq
+
+ tmpdir = pjoin(self.tmp_path, 'multi-parquet-uri-' + guid())
+
+ self.hdfs.mkdir(tmpdir)
+
+ expected = self._write_multiple_hdfs_pq_files(tmpdir)
+ path = _get_hdfs_uri(tmpdir)
+ result = pq.read_table(path)
+
+ _pandas_api.assert_frame_equal(result.to_pandas()
+ .sort_values(by='index')
+ .reset_index(drop=True),
+ expected.to_pandas())
+
+ @pytest.mark.pandas
+ @pytest.mark.parquet
+ def test_read_write_parquet_files_with_uri(self):
+ import pyarrow.parquet as pq
+
+ tmpdir = pjoin(self.tmp_path, 'uri-parquet-' + guid())
+ self.hdfs.mkdir(tmpdir)
+ path = _get_hdfs_uri(pjoin(tmpdir, 'test.parquet'))
+
+ size = 5
+ df = _test_dataframe(size, seed=0)
+ # Hack so that we don't have a dtype cast in v1 files
+ df['uint32'] = df['uint32'].astype(np.int64)
+ table = pa.Table.from_pandas(df, preserve_index=False)
+
+ pq.write_table(table, path, filesystem=self.hdfs)
+
+ result = pq.read_table(
+ path, filesystem=self.hdfs, use_legacy_dataset=True
+ ).to_pandas()
+
+ _pandas_api.assert_frame_equal(result, df)
+
+ @pytest.mark.parquet
+ @pytest.mark.pandas
+ def test_read_common_metadata_files(self):
+ tmpdir = pjoin(self.tmp_path, 'common-metadata-' + guid())
+ self.hdfs.mkdir(tmpdir)
+ _test_read_common_metadata_files(self.hdfs, tmpdir)
+
+ @pytest.mark.parquet
+ @pytest.mark.pandas
+ def test_write_to_dataset_with_partitions(self):
+ tmpdir = pjoin(self.tmp_path, 'write-partitions-' + guid())
+ self.hdfs.mkdir(tmpdir)
+ _test_write_to_dataset_with_partitions(
+ tmpdir, filesystem=self.hdfs)
+
+ @pytest.mark.parquet
+ @pytest.mark.pandas
+ def test_write_to_dataset_no_partitions(self):
+ tmpdir = pjoin(self.tmp_path, 'write-no_partitions-' + guid())
+ self.hdfs.mkdir(tmpdir)
+ _test_write_to_dataset_no_partitions(
+ tmpdir, filesystem=self.hdfs)
+
+
+class TestLibHdfs(HdfsTestCases, unittest.TestCase):
+
+ @classmethod
+ def check_driver(cls):
+ check_libhdfs_present()
+
+ def test_orphaned_file(self):
+ hdfs = hdfs_test_client()
+ file_path = self._make_test_file(hdfs, 'orphaned_file_test', 'fname',
+ b'foobarbaz')
+
+ f = hdfs.open(file_path)
+ hdfs = None
+ f = None # noqa
+
+
+def _get_hdfs_uri(path):
+ host = os.environ.get('ARROW_HDFS_TEST_HOST', 'localhost')
+ try:
+ port = int(os.environ.get('ARROW_HDFS_TEST_PORT', 0))
+ except ValueError:
+ raise ValueError('Env variable ARROW_HDFS_TEST_PORT was not '
+ 'an integer')
+ uri = "hdfs://{}:{}{}".format(host, port, path)
+
+ return uri
+
+
+@pytest.mark.hdfs
+@pytest.mark.pandas
+@pytest.mark.parquet
+@pytest.mark.fastparquet
+def test_fastparquet_read_with_hdfs():
+ from pandas.testing import assert_frame_equal
+
+ check_libhdfs_present()
+ try:
+ import snappy # noqa
+ except ImportError:
+ pytest.skip('fastparquet test requires snappy')
+
+ import pyarrow.parquet as pq
+ fastparquet = pytest.importorskip('fastparquet')
+
+ fs = hdfs_test_client()
+
+ df = util.make_dataframe()
+
+ table = pa.Table.from_pandas(df)
+
+ path = '/tmp/testing.parquet'
+ with fs.open(path, 'wb') as f:
+ pq.write_table(table, f)
+
+ parquet_file = fastparquet.ParquetFile(path, open_with=fs.open)
+
+ result = parquet_file.to_pandas()
+ assert_frame_equal(result, df)
projects / ceph.git / blobdiff