code: # Copyright 1999-2020 Alibaba Group Holding Ltd.
  #
  # Licensed 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 tempfile
  import unittest
  
  import numpy as np
  import pandas as pd
  
  import mars.dataframe as md
  from mars.config import option_context
  from mars.dataframe import DataFrame
  from mars.deploy.local.core import new_cluster
  from mars.session import new_session
  from mars.tests.core import TestBase, flaky
  
  try:
      import vineyard
  except ImportError:
      vineyard = None
  try:
      import sqlalchemy
  except ImportError:
      sqlalchemy = None
  try:
      import pyarrow as pa
  except ImportError:
      pa = None
  try:
      import fastparquet
  except ImportError:
      fastparquet = None
  
  _exec_timeout = 120 if 'CI' in os.environ else -1
  
  
  class Test(TestBase):
      def setUp(self):
          super().setUp()
          self.ctx, self.executor = self._create_test_context()
  
      def testToCSVExecution(self):
          index = pd.RangeIndex(100, 0, -1, name='index')
          raw = pd.DataFrame({
              'col1': np.random.rand(100),
              'col2': np.random.choice(['a', 'b', 'c'], (100,)),
              'col3': np.arange(100)
          }, index=index)
          df = DataFrame(raw, chunk_size=33)
  
          with tempfile.TemporaryDirectory() as base_path:
              # DATAFRAME TESTS
              # test one file with dataframe
              path = os.path.join(base_path, 'out.csv')
  
              r = df.to_csv(path)
              self.executor.execute_dataframe(r)
  
              result = pd.read_csv(path, dtype=raw.dtypes.to_dict())
              result.set_index('index', inplace=True)
              pd.testing.assert_frame_equal(result, raw)
  
              # test multi files with dataframe
              path = os.path.join(base_path, 'out-*.csv')
              r = df.to_csv(path)
              self.executor.execute_dataframe(r)
  
              dfs = [pd.read_csv(os.path.join(base_path, f'out-{i}.csv'),
                                 dtype=raw.dtypes.to_dict())
                     for i in range(4)]
              result = pd.concat(dfs, axis=0)
              result.set_index('index', inplace=True)
              pd.testing.assert_frame_equal(result, raw)
              pd.testing.assert_frame_equal(dfs[1].set_index('index'), raw.iloc[33: 66])
  
              with self.ctx:
                  # test df with unknown shape
                  df2 = DataFrame(raw, chunk_size=(50, 2))
                  df2 = df2[df2['col1'] < 1]
                  path2 = os.path.join(base_path, 'out2.csv')
                  r = df2.to_csv(path2)
                  self.executor.execute_dataframes([r])
  
                  result = pd.read_csv(path2, dtype=raw.dtypes.to_dict())
                  result.set_index('index', inplace=True)
                  pd.testing.assert_frame_equal(result, raw)
  
              # SERIES TESTS
              series = md.Series(raw.col1, chunk_size=33)
  
              # test one file with series
              path = os.path.join(base_path, 'out.csv')
              r = series.to_csv(path)
              self.executor.execute_dataframe(r)
  
              result = pd.read_csv(path, dtype=raw.dtypes.to_dict())
              result.set_index('index', inplace=True)
              pd.testing.assert_frame_equal(result, raw.col1.to_frame())
  
              # test multi files with series
              path = os.path.join(base_path, 'out-*.csv')
              r = series.to_csv(path)
              self.executor.execute_dataframe(r)
  
              dfs = [pd.read_csv(os.path.join(base_path, f'out-{i}.csv'),
                                 dtype=raw.dtypes.to_dict())
                     for i in range(4)]
              result = pd.concat(dfs, axis=0)
              result.set_index('index', inplace=True)
              pd.testing.assert_frame_equal(result, raw.col1.to_frame())
              pd.testing.assert_frame_equal(dfs[1].set_index('index'), raw.col1.to_frame().iloc[33: 66])
  
      @unittest.skipIf(sqlalchemy is None, 'sqlalchemy not installed')
      def testToSQL(self):
          index = pd.RangeIndex(100, 0, -1, name='index')
          raw = pd.DataFrame({
              'col1': np.random.rand(100),
              'col2': np.random.choice(['a', 'b', 'c'], (100,)),
              'col3': np.arange(100).astype('int64'),
          }, index=index)
  
          with tempfile.TemporaryDirectory() as d:
              table_name1 = 'test_table'
              table_name2 = 'test_table2'
              uri = 'sqlite:///' + os.path.join(d, 'test.db')
  
              engine = sqlalchemy.create_engine(uri)
  
              # test write dataframe
              df = DataFrame(raw, chunk_size=33)
              r = df.to_sql(table_name1, con=engine)
              self.executor.execute_dataframe(r)
  
              written = pd.read_sql(table_name1, con=engine, index_col='index') \
                  .sort_index(ascending=False)
              pd.testing.assert_frame_equal(raw, written)
  
              # test write with existing table
              with self.assertRaises(ValueError):
                  df.to_sql(table_name1, con=uri).execute()
  
              # test write series
              series = md.Series(raw.col1, chunk_size=33)
              with engine.connect() as conn:
                  r = series.to_sql(table_name2, con=conn)
                  self.executor.execute_dataframe(r)
  
              written = pd.read_sql(table_name2, con=engine, index_col='index') \
                  .sort_index(ascending=False)
              pd.testing.assert_frame_equal(raw.col1.to_frame(), written)
  
      @unittest.skipIf(vineyard is None, 'vineyard not installed')
      @flaky(max_runs=3)
      def testToVineyard(self):
          def run_with_given_session(session, **kw):
              ipc_socket = os.environ.get('VINEYARD_IPC_SOCKET', '/tmp/vineyard/vineyard.sock')
              with option_context({'vineyard.socket': ipc_socket}):
                  df1 = DataFrame(pd.DataFrame(np.arange(12).reshape(3, 4), columns=['a', 'b', 'c', 'd']),
                                  chunk_size=2)
                  object_id = df1.to_vineyard().execute(session=session, **kw).fetch(session=session)
                  df2 = md.from_vineyard(object_id)
  
                  df1_value = df1.execute(session=session, **kw).fetch(session=session)
                  df2_value = df2.execute(session=session, **kw).fetch(session=session)
                  pd.testing.assert_frame_equal(
                      df1_value.reset_index(drop=True), df2_value.reset_index(drop=True))
  
          with new_session().as_default() as session:
              run_with_given_session(session)
  
          with new_cluster(scheduler_n_process=2, worker_n_process=2,
                           shared_memory='20M', web=False) as cluster:
              with new_session(cluster.endpoint).as_default() as session:
                  run_with_given_session(session, timeout=_exec_timeout)
  
      @unittest.skipIf(pa is None, 'pyarrow not installed')
      def testToParquetArrowExecution(self):
          raw = pd.DataFrame({
              'col1': np.random.rand(100),
              'col2': np.arange(100),
              'col3': np.random.choice(['a', 'b', 'c'], (100,)),
          })
          df = DataFrame(raw, chunk_size=33)
  
          with tempfile.TemporaryDirectory() as base_path:
              # DATAFRAME TESTS
              path = os.path.join(base_path, 'out-*.parquet')
              r = df.to_parquet(path)
              self.executor.execute_dataframe(r)
  
              read_df = md.read_parquet(path)
              result = self.executor.execute_dataframe(read_df, concat=True)[0]
              result = result.sort_index()
              pd.testing.assert_frame_equal(result, raw)
  
              read_df = md.read_parquet(path)
              result = self.executor.execute_dataframe(read_df, concat=True)[0]
              result = result.sort_index()
              pd.testing.assert_frame_equal(result, raw)
  
              # test read_parquet then to_parquet
              read_df = md.read_parquet(path)
              r = read_df.to_parquet(path)
              self.executor.execute_dataframes([r])
  
              # test partition_cols
              path = os.path.join(base_path, 'out-partitioned')
              r = df.to_parquet(path, partition_cols=['col3'])
              self.executor.execute_dataframe(r)
  
              read_df = md.read_parquet(path)
              result = self.executor.execute_dataframe(read_df, concat=True)[0]
              result['col3'] = result['col3'].astype('object')
              pd.testing.assert_frame_equal(result.sort_values('col1').reset_index(drop=True),
                                            raw.sort_values('col1').reset_index(drop=True))
  
      @unittest.skipIf(fastparquet is None, 'fastparquet not installed')
      def testToParquetFastParquetExecution(self):
          raw = pd.DataFrame({
              'col1': np.random.rand(100),
              'col2': np.arange(100),
              'col3': np.random.choice(['a', 'b', 'c'], (100,)),
          })
          df = DataFrame(raw, chunk_size=33)
  
          with tempfile.TemporaryDirectory() as base_path:
              # test fastparquet
              path = os.path.join(base_path, 'out-fastparquet-*.parquet')
              r = df.to_parquet(path, engine='fastparquet', compression='gzip')
              self.executor.execute_dataframe(r)
  
apis: pandas.read_sql numpy.random.choice numpy.arange pandas.RangeIndex numpy.random.rand pandas.concat pandas.testing.assert_frame_equal