code: def test_optimization(setup):
      import sqlalchemy as sa
  
      with tempfile.TemporaryDirectory() as tempdir:
          filename = os.path.join(tempdir, "test_head.csv")
          rs = np.random.RandomState(0)
          pd_df = pd.DataFrame(
              {
                  "a": rs.randint(1000, size=(2000,)).astype(np.int64),
                  "b": rs.randint(1000, size=(2000,)).astype(np.int64),
                  "c": ["sss" for _ in range(2000)],
                  "d": ["eeee" for _ in range(2000)],
              }
          )
          pd_df.to_csv(filename, index=False)
  
          size = os.path.getsize(filename)
          chunk_bytes = size / 3 - 2
  
          df = md.read_csv(filename, chunk_bytes=chunk_bytes)
  
          cols = ["b", "a", "c"]
          r = df[cols]
          operand_executors = {
              DataFrameReadCSV: _wrap_execute_data_source_usecols(
                  cols, DataFrameReadCSV
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df[cols]
          result.reset_index(drop=True, inplace=True)
          pd.testing.assert_frame_equal(result, expected)
  
          cols = ["b", "a", "b"]
          r = df[cols].head(20)
          operand_executors = {
              DataFrameReadCSV: _wrap_execute_data_source_usecols(
                  cols, DataFrameReadCSV
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df[cols].head(20)
          result.reset_index(drop=True, inplace=True)
          pd.testing.assert_frame_equal(result, expected)
  
          r = df["c"]
          operand_executors = {
              DataFrameReadCSV: _wrap_execute_data_source_usecols(
                  "c", DataFrameReadCSV
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df["c"]
          result.reset_index(drop=True, inplace=True)
          pd.testing.assert_series_equal(result, expected)
  
          r = df["d"].head(3)
          operand_executors = {
              DataFrameReadCSV: _wrap_execute_data_source_mixed(
                  3, "d", DataFrameReadCSV
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df["d"].head(3)
          pd.testing.assert_series_equal(result, expected)
  
          # test DataFrame.head
          r = df.head(3)
          operand_executors = {
              DataFrameReadCSV: _wrap_execute_data_source(
                  3, DataFrameReadCSV
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df.head(3)
          pd.testing.assert_frame_equal(result, expected)
  
          # test DataFrame.tail
          r = df.tail(3)
          result = r.execute().fetch()
          expected = pd_df.tail(3)
          pd.testing.assert_frame_equal(
              result.reset_index(drop=True),
              expected.reset_index(drop=True),
          )
  
          # test head more than 1 chunk
          r = df.head(99)
          result = r.execute().fetch()
          result.reset_index(drop=True, inplace=True)
          expected = pd_df.head(99)
          pd.testing.assert_frame_equal(result, expected)
  
          # test Series.tail more than 1 chunk
          r = df.tail(99)
          result = r.execute().fetch()
          expected = pd_df.tail(99)
          pd.testing.assert_frame_equal(
              result.reset_index(drop=True),
              expected.reset_index(drop=True),
          )
  
          # test head number greater than limit
          df = md.read_csv(filename, chunk_bytes=chunk_bytes)
          r = df.head(1100)
  
          with pytest.raises(RuntimeError):
              operand_executors = {
                  DataFrameReadCSV: _wrap_execute_data_source(
                      3, DataFrameReadCSV
                  )
              }
              r.execute(
                  extra_config={"operand_executors": operand_executors}
              )
  
          result = r.execute().fetch()
          expected = pd_df.head(1100)
          pd.testing.assert_frame_equal(
              result.reset_index(drop=True),
              expected.reset_index(drop=True),
          )
  
          filename = os.path.join(tempdir, "test_sql.db")
          conn = sa.create_engine("sqlite:///" + filename)
          pd_df.to_sql("test_sql", conn)
  
          df = md.read_sql(
              "test_sql", conn, index_col="index", chunk_size=20
          )
  
          # test DataFrame.head
          r = df.head(3)
          operand_executors = {
              DataFrameReadSQL: _wrap_execute_data_source(
                  3, DataFrameReadSQL
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          result.index.name = None
          expected = pd_df.head(3)
          pd.testing.assert_frame_equal(result, expected)
  
          # test head on read_parquet
          filename = os.path.join(tempdir, "test_parquet.db")
          pd_df.to_parquet(filename, index=False, compression="gzip")
  
          engines = []
          if pa is not None:
              engines.append("pyarrow")
          if fp is not None:
              engines.append("fastparquet")
  
          for engine in engines:
              df = md.read_parquet(filename, engine=engine)
              r = df.head(3)
  
              operand_executors = {
                  DataFrameReadParquet: _wrap_execute_data_source(
                      3, DataFrameReadParquet
                  )
              }
              result = r.execute(
                  extra_config={"operand_executors": operand_executors}
              ).fetch()
              expected = pd_df.head(3)
              pd.testing.assert_frame_equal(result, expected)
  
          dirname = os.path.join(tempdir, "test_parquet2")
          os.makedirs(dirname)
          pd_df[:1000].to_parquet(os.path.join(dirname, "q1.parquet"))
          pd_df[1000:].to_parquet(os.path.join(dirname, "q2.parquet"))
  
          df = md.read_parquet(dirname)
          r = df.head(3)
  
          operand_executors = {
              DataFrameReadParquet: _wrap_execute_data_source(
                  3, DataFrameReadParquet
              )
          }
          result = r.execute(
              extra_config={"operand_executors": operand_executors}
          ).fetch()
          expected = pd_df.head(3)
          pd.testing.assert_frame_equal(result, expected)
  
imports: import os
  import tempfile
  import mars
  import numpy
  import pandas
  import pytest
  from ....None import dataframe
  from ....None import tensor
  from ....utils import pd_release_version
  from ...datasource.read_csv import DataFrameReadCSV
  from ...datasource.read_sql import DataFrameReadSQL
  from ...datasource.read_parquet import DataFrameReadParquet
label_desc: Use image node_id 21 for calling a global function with example usage: test_optimization(setup) without return types