text: # 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 datetime
  from collections import OrderedDict
  
  import numpy as np
  import pytest
  
  import pyarrow as pa
  from pyarrow.tests.parquet.common import _check_roundtrip, make_sample_file
  
  try:
      import pyarrow.parquet as pq
      from pyarrow.tests.parquet.common import _write_table
  except ImportError:
      pq = None
  
  
  try:
      import pandas as pd
      import pandas.testing as tm
  
      from pyarrow.tests.parquet.common import alltypes_sample
  except ImportError:
      pd = tm = None
  
  
  pytestmark = pytest.mark.parquet
  
  
  @pytest.mark.pandas
  def test_parquet_metadata_api():
      df = alltypes_sample(size=10000)
      df = df.reindex(columns=sorted(df.columns))
      df.index = np.random.randint(0, 1000000, size=len(df))
  
      fileh = make_sample_file(df)
      ncols = len(df.columns)
  
      # Series of sniff tests
      meta = fileh.metadata
      repr(meta)
      assert meta.num_rows == len(df)
      assert meta.num_columns == ncols + 1  # +1 for index
      assert meta.num_row_groups == 1
      assert meta.format_version == '2.0'
      assert 'parquet-cpp' in meta.created_by
      assert isinstance(meta.serialized_size, int)
      assert isinstance(meta.metadata, dict)
  
      # Schema
      schema = fileh.schema
      assert meta.schema is schema
      assert len(schema) == ncols + 1  # +1 for index
      repr(schema)
  
      col = schema[0]
      repr(col)
      assert col.name == df.columns[0]
      assert col.max_definition_level == 1
      assert col.max_repetition_level == 0
      assert col.max_repetition_level == 0
  
      assert col.physical_type == 'BOOLEAN'
      assert col.converted_type == 'NONE'
  
      with pytest.raises(IndexError):
          schema[ncols + 1]  # +1 for index
  
      with pytest.raises(IndexError):
          schema[-1]
  
      # Row group
      for rg in range(meta.num_row_groups):
          rg_meta = meta.row_group(rg)
          assert isinstance(rg_meta, pq.RowGroupMetaData)
          repr(rg_meta)
  
          for col in range(rg_meta.num_columns):
              col_meta = rg_meta.column(col)
              assert isinstance(col_meta, pq.ColumnChunkMetaData)
              repr(col_meta)
  
      with pytest.raises(IndexError):
          meta.row_group(-1)
  
      with pytest.raises(IndexError):
          meta.row_group(meta.num_row_groups + 1)
  
      rg_meta = meta.row_group(0)
      assert rg_meta.num_rows == len(df)
      assert rg_meta.num_columns == ncols + 1  # +1 for index
      assert rg_meta.total_byte_size > 0
  
      with pytest.raises(IndexError):
          col_meta = rg_meta.column(-1)
  
      with pytest.raises(IndexError):
          col_meta = rg_meta.column(ncols + 2)
  
      col_meta = rg_meta.column(0)
      assert col_meta.file_offset > 0
      assert col_meta.file_path == ''  # created from BytesIO
      assert col_meta.physical_type == 'BOOLEAN'
      assert col_meta.num_values == 10000
      assert col_meta.path_in_schema == 'bool'
      assert col_meta.is_stats_set is True
      assert isinstance(col_meta.statistics, pq.Statistics)
      assert col_meta.compression == 'SNAPPY'
      assert col_meta.encodings == ('PLAIN', 'RLE')
      assert col_meta.has_dictionary_page is False
      assert col_meta.dictionary_page_offset is None
      assert col_meta.data_page_offset > 0
      assert col_meta.total_compressed_size > 0
      assert col_meta.total_uncompressed_size > 0
      with pytest.raises(NotImplementedError):
          col_meta.has_index_page
      with pytest.raises(NotImplementedError):
          col_meta.index_page_offset
  
  
  def test_parquet_metadata_lifetime(tempdir):
      # ARROW-6642 - ensure that chained access keeps parent objects alive
      table = pa.table({'a': [1, 2, 3]})
      pq.write_table(table, tempdir / 'test_metadata_segfault.parquet')
      dataset = pq.ParquetDataset(tempdir / 'test_metadata_segfault.parquet')
      dataset.pieces[0].get_metadata().row_group(0).column(0).statistics
  
  
  @pytest.mark.pandas
  @pytest.mark.parametrize(
      (
          'data',
          'type',
          'physical_type',
          'min_value',
          'max_value',
          'null_count',
          'num_values',
          'distinct_count'
      ),
      [
          ([1, 2, 2, None, 4], pa.uint8(), 'INT32', 1, 4, 1, 4, 0),
          ([1, 2, 2, None, 4], pa.uint16(), 'INT32', 1, 4, 1, 4, 0),
          ([1, 2, 2, None, 4], pa.uint32(), 'INT32', 1, 4, 1, 4, 0),
          ([1, 2, 2, None, 4], pa.uint64(), 'INT64', 1, 4, 1, 4, 0),
          ([-1, 2, 2, None, 4], pa.int8(), 'INT32', -1, 4, 1, 4, 0),
          ([-1, 2, 2, None, 4], pa.int16(), 'INT32', -1, 4, 1, 4, 0),
          ([-1, 2, 2, None, 4], pa.int32(), 'INT32', -1, 4, 1, 4, 0),
          ([-1, 2, 2, None, 4], pa.int64(), 'INT64', -1, 4, 1, 4, 0),
          (
              [-1.1, 2.2, 2.3, None, 4.4], pa.float32(),
              'FLOAT', -1.1, 4.4, 1, 4, 0
          ),
          (
              [-1.1, 2.2, 2.3, None, 4.4], pa.float64(),
              'DOUBLE', -1.1, 4.4, 1, 4, 0
          ),
          (
              ['', 'b', chr(1000), None, 'aaa'], pa.binary(),
              'BYTE_ARRAY', b'', chr(1000).encode('utf-8'), 1, 4, 0
          ),
          (
              [True, False, False, True, True], pa.bool_(),
              'BOOLEAN', False, True, 0, 5, 0
          ),
          (
              [b'\x00', b'b', b'12', None, b'aaa'], pa.binary(),
              'BYTE_ARRAY', b'\x00', b'b', 1, 4, 0
          ),
      ]
  )
  def test_parquet_column_statistics_api(data, type, physical_type, min_value,
                                         max_value, null_count, num_values,
                                         distinct_count):
      df = pd.DataFrame({'data': data})
      schema = pa.schema([pa.field('data', type)])
      table = pa.Table.from_pandas(df, schema=schema, safe=False)
      fileh = make_sample_file(table)
  
      meta = fileh.metadata
  
      rg_meta = meta.row_group(0)
      col_meta = rg_meta.column(0)
  
      stat = col_meta.statistics
      assert stat.has_min_max
      assert _close(type, stat.min, min_value)
      assert _close(type, stat.max, max_value)
      assert stat.null_count == null_count
      assert stat.num_values == num_values
      # TODO(kszucs) until parquet-cpp API doesn't expose HasDistinctCount
      # method, missing distinct_count is represented as zero instead of None
      assert stat.distinct_count == distinct_count
      assert stat.physical_type == physical_type
  
  
  def _close(type, left, right):
      if type == pa.float32():
          return abs(left - right) < 1E-7
      elif type == pa.float64():
          return abs(left - right) < 1E-13
      else:
          return left == right
  
  
  # ARROW-6339
  @pytest.mark.pandas
  def test_parquet_raise_on_unset_statistics():
      df = pd.DataFrame({"t": pd.Series([pd.NaT], dtype="datetime64[ns]")})
      meta = make_sample_file(pa.Table.from_pandas(df)).metadata
  
      assert not meta.row_group(0).column(0).statistics.has_min_max
      assert meta.row_group(0).column(0).statistics.max is None
  
  
  def test_statistics_convert_logical_types(tempdir):
      # ARROW-5166, ARROW-4139
  
      # (min, max, type)
      cases = [(10, 11164359321221007157, pa.uint64()),
               (10, 4294967295, pa.uint32()),
               ("ähnlich", "öffentlich", pa.utf8()),
               (datetime.time(10, 30, 0, 1000), datetime.time(15, 30, 0, 1000),
                pa.time32('ms')),
               (datetime.time(10, 30, 0, 1000), datetime.time(15, 30, 0, 1000),
                pa.time64('us')),
               (datetime.datetime(2019, 6, 24, 0, 0, 0, 1000),
                datetime.datetime(2019, 6, 25, 0, 0, 0, 1000),
                pa.timestamp('ms')),
               (datetime.datetime(2019, 6, 24, 0, 0, 0, 1000),
                datetime.datetime(2019, 6, 25, 0, 0, 0, 1000),
                pa.timestamp('us'))]
  
      for i, (min_val, max_val, typ) in enumerate(cases):
          t = pa.Table.from_arrays([pa.array([min_val, max_val], type=typ)],
                                   ['col'])
          path = str(tempdir / ('example{}.parquet'.format(i)))
          pq.write_table(t, path, version='2.0')
          pf = pq.ParquetFile(path)
          stats = pf.metadata.row_group(0).column(0).statistics
          assert stats.min == min_val
          assert stats.max == max_val
  
  
  def test_parquet_write_disable_statistics(tempdir):
      table = pa.Table.from_pydict(
          OrderedDict([
              ('a', pa.array([1, 2, 3])),
              ('b', pa.array(['a', 'b', 'c']))
          ])
      )
      _write_table(table, tempdir / 'data.parquet')
      meta = pq.read_metadata(tempdir / 'data.parquet')
      for col in [0, 1]:
          cc = meta.row_group(0).column(col)
          assert cc.is_stats_set is True
          assert cc.statistics is not None
  
      _write_table(table, tempdir / 'data2.parquet', write_statistics=False)
      meta = pq.read_metadata(tempdir / 'data2.parquet')
      for col in [0, 1]:
          cc = meta.row_group(0).column(col)
          assert cc.is_stats_set is False
          assert cc.statistics is None
  
      _write_table(table, tempdir / 'data3.parquet', write_statistics=['a'])
      meta = pq.read_metadata(tempdir / 'data3.parquet')
      cc_a = meta.row_group(0).column(0)
      cc_b = meta.row_group(0).column(1)
      assert cc_a.is_stats_set is True
      assert cc_b.is_stats_set is False
      assert cc_a.statistics is not None
      assert cc_b.statistics is None
  
  
  def test_field_id_metadata():
      # ARROW-7080
      table = pa.table([pa.array([1], type='int32'),
                        pa.array([[]], type=pa.list_(pa.int32())),
                        pa.array([b'boo'], type='binary')],
                       ['f0', 'f1', 'f2'])
  
      bio = pa.BufferOutputStream()
      pq.write_table(table, bio)
      contents = bio.getvalue()
  
      pf = pq.ParquetFile(pa.BufferReader(contents))
      schema = pf.schema_arrow
  
      # Expected Parquet schema for reference
      #
      # required group field_id=0 schema {
      #   optional int32 field_id=1 f0;
      #   optional group field_id=2 f1 (List) {
      #     repeated group field_id=3 list {
      #       optional int32 field_id=4 item;
      #     }
      #   }
      #   optional binary field_id=5 f2;
      # }
  
      field_name = b'PARQUET:field_id'
      assert schema[0].metadata[field_name] == b'1'
  
      list_field = schema[1]
      assert list_field.metadata[field_name] == b'2'
  
      list_item_field = list_field.type.value_field
      assert list_item_field.metadata[field_name] == b'4'
  
      assert schema[2].metadata[field_name] == b'5'
  
  
  @pytest.mark.pandas
  def test_multi_dataset_metadata(tempdir):
      filenames = ["ARROW-1983-dataset.0", "ARROW-1983-dataset.1"]
      metapath = str(tempdir / "_metadata")
  
      # create a test dataset
      df = pd.DataFrame({
          'one': [1, 2, 3],
          'two': [-1, -2, -3],
          'three': [[1, 2], [2, 3], [3, 4]],
      })
      table = pa.Table.from_pandas(df)
  
      # write dataset twice and collect/merge metadata
      _meta = None
      for filename in filenames:
          meta = []
          pq.write_table(table, str(tempdir / filename),
                         metadata_collector=meta)
          meta[0].set_file_path(filename)
          if _meta is None:
              _meta = meta[0]
          else:
              _meta.append_row_groups(meta[0])
  
      # Write merged metadata-only file
      with open(metapath, "wb") as f:
          _meta.write_metadata_file(f)
  
      # Read back the metadata
      meta = pq.read_metadata(metapath)
      md = meta.to_dict()
      _md = _meta.to_dict()
      for key in _md:
          if key != 'serialized_size':
              assert _md[key] == md[key]
      assert _md['num_columns'] == 3
      assert _md['num_rows'] == 6
      assert _md['num_row_groups'] == 2
      assert _md['serialized_size'] == 0
      assert md['serialized_size'] > 0
  
  
  def test_write_metadata(tempdir):
      path = str(tempdir / "metadata")
      schema = pa.schema([("a", "int64"), ("b", "float64")])
  
      # write a pyarrow schema
      pq.write_metadata(schema, path)
      parquet_meta = pq.read_metadata(path)
      schema_as_arrow = parquet_meta.schema.to_arrow_schema()
      assert schema_as_arrow.equals(schema)
  
      # ARROW-8980: Check that the ARROW:schema metadata key was removed
      if schema_as_arrow.metadata:
          assert b'ARROW:schema' not in schema_as_arrow.metadata
  
      # pass through writer keyword arguments
      for version in ["1.0", "2.0"]:
          pq.write_metadata(schema, path, version=version)
          parquet_meta = pq.read_metadata(path)
          assert parquet_meta.format_version == version
  
      # metadata_collector: list of FileMetaData objects
      table = pa.table({'a': [1, 2], 'b': [.1, .2]}, schema=schema)
      pq.write_table(table, tempdir / "data.parquet")
      parquet_meta = pq.read_metadata(str(tempdir / "data.parquet"))
      pq.write_metadata(
          schema, path, metadata_collector=[parquet_meta, parquet_meta]
      )
      parquet_meta_mult = pq.read_metadata(path)
      assert parquet_meta_mult.num_row_groups == 2
  
      # append metadata with different schema raises an error
      with pytest.raises(RuntimeError, match="requires equal schemas"):
          pq.write_metadata(
              pa.schema([("a", "int32"), ("b", "null")]),
              path, metadata_collector=[parquet_meta, parquet_meta]
          )
  
  
  def test_table_large_metadata():
      # ARROW-8694
      my_schema = pa.schema([pa.field('f0', 'double')],
                            metadata={'large': 'x' * 10000000})
  
      table = pa.table([np.arange(10)], schema=my_schema)
      _check_roundtrip(table)
  
  
  @pytest.mark.pandas
  def test_compare_schemas():
      df = alltypes_sample(size=10000)
  
      fileh = make_sample_file(df)
      fileh2 = make_sample_file(df)
      fileh3 = make_sample_file(df[df.columns[::2]])
  
      # ParquetSchema
      assert isinstance(fileh.schema, pq.ParquetSchema)
      assert fileh.schema.equals(fileh.schema)
      assert fileh.schema == fileh.schema
      assert fileh.schema.equals(fileh2.schema)
      assert fileh.schema == fileh2.schema
      assert fileh.schema != 'arbitrary object'
      assert not fileh.schema.equals(fileh3.schema)
      assert fileh.schema != fileh3.schema
  
      # ColumnSchema
      assert isinstance(fileh.schema[0], pq.ColumnSchema)
      assert fileh.schema[0].equals(fileh.schema[0])
      assert fileh.schema[0] == fileh.schema[0]
      assert not fileh.schema[0].equals(fileh.schema[1])
      assert fileh.schema[0] != fileh.schema[1]
      assert fileh.schema[0] != 'arbitrary object'
  
  
  @pytest.mark.pandas
  def test_read_schema(tempdir):
      N = 100
      df = pd.DataFrame({
          'index': np.arange(N),
          'values': np.random.randn(N)
      }, columns=['index', 'values'])
  
      data_path = tempdir / 'test.parquet'
  
      table = pa.Table.from_pandas(df)
      _write_table(table, data_path)
  
      read1 = pq.read_schema(data_path)
      read2 = pq.read_schema(data_path, memory_map=True)
      assert table.schema.equals(read1)
      assert table.schema.equals(read2)
  
      assert table.schema.metadata[b'pandas'] == read1.metadata[b'pandas']
  
  
  def test_parquet_metadata_empty_to_dict(tempdir):
      # https://issues.apache.org/jira/browse/ARROW-10146
      table = pa.table({"a": pa.array([], type="int64")})
      pq.write_table(table, tempdir / "data.parquet")
      metadata = pq.read_metadata(tempdir / "data.parquet")
      # ensure this doesn't error / statistics set to None
      metadata_dict = metadata.to_dict()
      assert len(metadata_dict["row_groups"]) == 1
      assert len(metadata_dict["row_groups"][0]["columns"]) == 1
      assert metadata_dict["row_groups"][0]["columns"][0]["statistics"] is None
  