#

# 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.

#

"""Tests for schemas."""

# pytype: skip-file

import typing

import unittest

import numpy as np

import pandas as pd

from parameterized import parameterized

import apache_beam as beam

from apache_beam.coders import RowCoder

from apache_beam.coders.typecoders import registry as coders_registry

from apache_beam.dataframe import schemas

from apache_beam.dataframe import transforms

from apache_beam.testing.test_pipeline import TestPipeline

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

from apache_beam.typehints import row_type

from apache_beam.typehints import typehints

from apache_beam.typehints.native_type_compatibility import match_is_named_tuple

Simple = typing.NamedTuple(

'Simple', [('name', str), ('id', int), ('height', float)])

coders_registry.register_coder(Simple, RowCoder)

Animal = typing.NamedTuple(

'Animal', [('animal', str), ('max_speed', typing.Optional[float])])

coders_registry.register_coder(Animal, RowCoder)

def matches_df(expected):

def check_df_pcoll_equal(actual):

actual = pd.concat(actual)

sorted_actual = actual.sort_values(by=list(actual.columns)).reset_index(

drop=True)

sorted_expected = expected.sort_values(

by=list(expected.columns)).reset_index(drop=True)

pd.testing.assert_frame_equal(sorted_actual, sorted_expected)

return check_df_pcoll_equal

# Test data for all supported types that can be easily tested.

# Excludes bytes because it's difficult to create a series and dataframe bytes

# dtype. For example:

# pd.Series([b'abc'], dtype=bytes).dtype != 'S'

# pd.Series([b'abc'], dtype=bytes).astype(bytes).dtype == 'S'

# (test data, pandas_type, column_name, beam_type)

COLUMNS: typing.List[typing.Tuple[typing.List[typing.Any],

typing.Any,

str,

typing.Any]] = [

([375, 24, 0, 10, 16],

np.int32,

'i32',

np.int32),

([375, 24, 0, 10, 16],

np.int64,

'i64',

np.int64),

([375, 24, None, 10, 16],

pd.Int32Dtype(),

'i32_nullable',

typing.Optional[np.int32]),

([375, 24, None, 10, 16],

pd.Int64Dtype(),

'i64_nullable',

typing.Optional[np.int64]),

([375., 24., None, 10., 16.],

np.float64,

'f64',

typing.Optional[np.float64]),

([375., 24., None, 10., 16.],

np.float32,

'f32',

typing.Optional[np.float32]),

([True, False, True, True, False],

bool,

'bool',

bool),

(['Falcon', 'Ostrich', None, 3.14, 0],

object,

'any',

typing.Any),

([True, False, True, None, False],

pd.BooleanDtype(),

'bool_nullable',

typing.Optional[bool]),

([

'Falcon',

'Ostrich',

None,

'Aardvark',

'Elephant'

],

pd.StringDtype(),

'strdtype',

typing.Optional[str]),

]

NICE_TYPES_DF = pd.DataFrame(columns=[name for _, _, name, _ in COLUMNS])

for arr, dtype, name, _ in COLUMNS:

NICE_TYPES_DF[name] = pd.Series(arr, dtype=dtype, name=name).astype(dtype)

NICE_TYPES_PROXY = NICE_TYPES_DF[:0]

SERIES_TESTS = [(pd.Series(arr, dtype=dtype, name=name), arr, beam_type)

for (arr, dtype, name, beam_type) in COLUMNS]

_TEST_ARRAYS: typing.List[typing.List[typing.Any]] = [

arr for (arr, _, _, _) in COLUMNS

]

DF_RESULT = list(zip(*_TEST_ARRAYS))

BEAM_SCHEMA = typing.NamedTuple( # type: ignore

'BEAM_SCHEMA', [(name, beam_type) for _, _, name, beam_type in COLUMNS])

INDEX_DF_TESTS = [(

NICE_TYPES_DF.set_index([name for _, _, name, _ in COLUMNS[:i]]),

DF_RESULT,

BEAM_SCHEMA) for i in range(1, len(COLUMNS) + 1)]

NOINDEX_DF_TESTS = [(NICE_TYPES_DF, DF_RESULT, BEAM_SCHEMA)]

# Get major, minor, bugfix version

PD_VERSION = tuple(map(int, pd.__version__.split('.')[0:3]))

def test_name_func(testcase_func, param_num, params):

df_or_series, _, _ = params.args

if isinstance(df_or_series, pd.Series):

return f"{testcase_func.__name__}_Series[{df_or_series.dtype}]"

elif isinstance(df_or_series, pd.DataFrame):

return (

f"{testcase_func.__name__}_DataFrame"

f"[{','.join(str(dtype) for dtype in df_or_series.dtypes)}]")

else:

raise ValueError(

f"Encountered unsupported param in {testcase_func.__name__}. "

"Expected Series or DataFrame, got:\n" + str(df_or_series))

class SchemasTest(unittest.TestCase):

def test_simple_df(self):

expected = pd.DataFrame({

'name': list(str(i) for i in range(5)),

'id': list(range(5)),

'height': list(float(i) for i in range(5))

},

columns=['name', 'id', 'height'])

expected.name = expected.name.astype(pd.StringDtype())

with TestPipeline() as p:

res = (

p

| beam.Create(

[Simple(name=str(i), id=i, height=float(i)) for i in range(5)])

| schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10))

assert_that(res, matches_df(expected))

def test_simple_df_with_beam_row(self):

expected = pd.DataFrame({

'name': list(str(i) for i in range(5)),

'id': list(range(5)),

'height': list(float(i) for i in range(5))

},

columns=['name', 'id', 'height'])

expected.name = expected.name.astype(pd.StringDtype())

with TestPipeline() as p:

res = (

p

| beam.Create([(str(i), i, float(i)) for i in range(5)])

| beam.Select(

name=lambda r: str(r[0]),

id=lambda r: int(r[1]),

height=lambda r: float(r[2]))

| schemas.BatchRowsAsDataFrame(min_batch_size=10, max_batch_size=10))

assert_that(res, matches_df(expected))

def test_generate_proxy(self):

expected = pd.DataFrame({

'animal': pd.Series(dtype=pd.StringDtype()),

'max_speed': pd.Series(dtype=np.float64)

})

pd.testing.assert_frame_equal(schemas.generate_proxy(Animal), expected)

def test_generate_proxy_beam_typehint(self):

expected = pd.Series(dtype=pd.Int32Dtype())

actual = schemas.generate_proxy(typehints.Optional[np.int32])

pd.testing.assert_series_equal(actual, expected)

def test_nice_types_proxy_roundtrip(self):

roundtripped = schemas.generate_proxy(

schemas.element_type_from_dataframe(NICE_TYPES_PROXY))

self.assertTrue(roundtripped.equals(NICE_TYPES_PROXY))

@unittest.skipIf(

PD_VERSION == (1, 2, 1),

"Can't roundtrip bytes in pandas 1.2.1"

"https://github.com/pandas-dev/pandas/issues/39474")

def test_bytes_proxy_roundtrip(self):

proxy = pd.DataFrame({'bytes': []})

proxy.bytes = proxy.bytes.astype(bytes)

roundtripped = schemas.generate_proxy(

schemas.element_type_from_dataframe(proxy))

self.assertEqual(roundtripped.bytes.dtype.kind, 'S')

def test_batch_with_df_transform(self):

with TestPipeline() as p:

res = (

p

| beam.Create([

Animal('Falcon', 380.0),

Animal('Falcon', 370.0),

Animal('Parrot', 24.0),

Animal('Parrot', 26.0)

])

| schemas.BatchRowsAsDataFrame()

| transforms.DataframeTransform(

lambda df: df.groupby('animal').mean(),

# TODO: Generate proxy in this case as well

proxy=schemas.generate_proxy(Animal),

include_indexes=True))

assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

# Do the same thing, but use reset_index() to make sure 'animal' is included

with TestPipeline() as p:

with beam.dataframe.allow_non_parallel_operations():

res = (

p

| beam.Create([

Animal('Falcon', 380.0),

Animal('Falcon', 370.0),

Animal('Parrot', 24.0),

Animal('Parrot', 26.0)

])

| schemas.BatchRowsAsDataFrame()

| transforms.DataframeTransform(

lambda df: df.groupby('animal').mean().reset_index(),

# TODO: Generate proxy in this case as well

proxy=schemas.generate_proxy(Animal)))

assert_that(res, equal_to([('Falcon', 375.), ('Parrot', 25.)]))

def assert_typehints_equal(self, left, right):

def maybe_drop_rowtypeconstraint(typehint):

if isinstance(typehint, row_type.RowTypeConstraint):

return typehint.user_type

else:

return typehint

left = maybe_drop_rowtypeconstraint(typehints.normalize(left))

right = maybe_drop_rowtypeconstraint(typehints.normalize(right))

if match_is_named_tuple(left):

self.assertTrue(match_is_named_tuple(right))

self.assertEqual(left.__annotations__, right.__annotations__)

else:

self.assertEqual(left, right)

@parameterized.expand(

SERIES_TESTS + NOINDEX_DF_TESTS, name_func=test_name_func)

def test_unbatch_no_index(self, df_or_series, rows, beam_type):

proxy = df_or_series[:0]

with TestPipeline() as p:

res = (

p | beam.Create([df_or_series[::2], df_or_series[1::2]])

| schemas.UnbatchPandas(proxy))

# Verify that the unbatched PCollection has the expected typehint

# TODO(https://github.com/apache/beam/issues/19923): typehints should

# support NamedTuple so we can use typehints.is_consistent_with here

# instead

self.assert_typehints_equal(res.element_type, beam_type)

assert_that(res, equal_to(rows))

@parameterized.expand(SERIES_TESTS + INDEX_DF_TESTS, name_func=test_name_func)

def test_unbatch_with_index(self, df_or_series, rows, _):

proxy = df_or_series[:0]

if (PD_VERSION < (1, 2) and

set(['i32_nullable', 'i64_nullable']).intersection(proxy.index.names)):

self.skipTest(

"pandas<1.2 incorrectly changes Int64Dtype to int64 when "

"moved to index.")

with TestPipeline() as p:

res = (

p | beam.Create([df_or_series[::2], df_or_series[1::2]])

| schemas.UnbatchPandas(proxy, include_indexes=True))

assert_that(res, equal_to(rows))

@parameterized.expand(SERIES_TESTS, name_func=test_name_func)

def test_unbatch_series_with_index_warns(

self, series, unused_rows, unused_type):

proxy = series[:0]

with TestPipeline() as p:

input_pc = p | beam.Create([series[::2], series[1::2]])

with self.assertWarns(UserWarning):

_ = input_pc | schemas.UnbatchPandas(proxy, include_indexes=True)

def test_unbatch_include_index_unnamed_index_raises(self):

df = pd.DataFrame({'foo': [1, 2, 3, 4]})

proxy = df[:0]

with TestPipeline() as p:

pc = p | beam.Create([df[::2], df[1::2]])

with self.assertRaisesRegex(ValueError, 'unnamed'):

_ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

def test_unbatch_include_index_nonunique_index_raises(self):

df = pd.DataFrame({'foo': [1, 2, 3, 4]})

df.index = pd.MultiIndex.from_arrays([[1, 2, 3, 4], [4, 3, 2, 1]],

names=['bar', 'bar'])

proxy = df[:0]

with TestPipeline() as p:

pc = p | beam.Create([df[::2], df[1::2]])

with self.assertRaisesRegex(ValueError, 'bar'):

_ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

def test_unbatch_include_index_column_conflict_raises(self):

df = pd.DataFrame({'foo': [1, 2, 3, 4]})

df.index = pd.Index([4, 3, 2, 1], name='foo')

proxy = df[:0]

with TestPipeline() as p:

pc = p | beam.Create([df[::2], df[1::2]])

with self.assertRaisesRegex(ValueError, 'foo'):

_ = pc | schemas.UnbatchPandas(proxy, include_indexes=True)

def test_unbatch_datetime(self):

s = pd.Series(

pd.date_range(

'1/1/2000', periods=100, freq='m', tz='America/Los_Angeles'))

proxy = s[:0]

with TestPipeline() as p:

res = (

p | beam.Create([s[::2], s[1::2]])

| schemas.UnbatchPandas(proxy, include_indexes=True))

assert_that(res, equal_to(list(s)))

if __name__ == '__main__':

unittest.main()