#

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

import pandas as pd

import apache_beam as beam

from apache_beam.dataframe import convert

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

def equal_to_unordered_series(expected):

def check(actual):

actual = pd.concat(actual)

if sorted(expected) != sorted(actual):

raise AssertionError('Series not equal: \n%s\n%s\n' % (expected, actual))

return check

class ConvertTest(unittest.TestCase):

def test_convert_yield_pandas(self):

with beam.Pipeline() as p:

a = pd.Series([1, 2, 3])

b = pd.Series([100, 200, 300])

pc_a = p | 'A' >> beam.Create([a])

pc_b = p | 'B' >> beam.Create([b])

df_a = convert.to_dataframe(pc_a, proxy=a[:0])

df_b = convert.to_dataframe(pc_b, proxy=b[:0])

df_2a = 2 * df_a

df_3a = 3 * df_a

df_ab = df_a * df_b

# Converting multiple results at a time can be more efficient.

pc_2a, pc_ab = convert.to_pcollection(df_2a, df_ab,

yield_elements='pandas')

# But separate conversions can be done as well.

pc_3a = convert.to_pcollection(df_3a, yield_elements='pandas')

assert_that(pc_2a, equal_to_unordered_series(2 * a), label='Check2a')

assert_that(pc_3a, equal_to_unordered_series(3 * a), label='Check3a')

assert_that(pc_ab, equal_to_unordered_series(a * b), label='Checkab')

def test_convert(self):

with beam.Pipeline() as p:

a = pd.Series([1, 2, 3])

b = pd.Series([100, 200, 300])

pc_a = p | 'A' >> beam.Create(a)

pc_b = p | 'B' >> beam.Create(b)

df_a = convert.to_dataframe(pc_a)

df_b = convert.to_dataframe(pc_b)

df_2a = 2 * df_a

df_3a = 3 * df_a

df_ab = df_a * df_b

# Converting multiple results at a time can be more efficient.

pc_2a, pc_ab = convert.to_pcollection(df_2a, df_ab)

# But separate conversions can be done as well.

pc_3a = convert.to_pcollection(df_3a)

assert_that(pc_2a, equal_to(list(2 * a)), label='Check2a')

assert_that(pc_3a, equal_to(list(3 * a)), label='Check3a')

assert_that(pc_ab, equal_to(list(a * b)), label='Checkab')

def test_convert_with_none(self):

# Ensure the logical Any type allows (nullable) None, see BEAM-12587.

df = pd.DataFrame({'A': ['str', 10, None], 'B': [None, 'str', 20]})

with beam.Pipeline() as p:

res = convert.to_pcollection(df, pipeline=p) | beam.Map(tuple)

assert_that(res, equal_to([(row.A, row.B) for _, row in df.iterrows()]))

def test_convert_scalar(self):

with beam.Pipeline() as p:

pc = p | 'A' >> beam.Create([1, 2, 3])

s = convert.to_dataframe(pc)

pc_sum = convert.to_pcollection(s.sum())

assert_that(pc_sum, equal_to([6]))

def test_convert_non_deferred(self):

with beam.Pipeline() as p:

s1 = pd.Series([1, 2, 3])

s2 = convert.to_dataframe(p | beam.Create([100, 200, 300]))

pc1, pc2 = convert.to_pcollection(s1, s2, pipeline=p)

assert_that(pc1, equal_to([1, 2, 3]), label='CheckNonDeferred')

assert_that(pc2, equal_to([100, 200, 300]), label='CheckDeferred')

def test_convert_memoization(self):

with beam.Pipeline() as p:

a = pd.Series([1, 2, 3])

b = pd.Series([100, 200, 300])

pc_a = p | 'A' >> beam.Create([a])

pc_b = p | 'B' >> beam.Create([b])

df_a = convert.to_dataframe(pc_a, proxy=a[:0])

df_b = convert.to_dataframe(pc_b, proxy=b[:0])

df_2a = 2 * df_a

df_3a = 3 * df_a

df_ab = df_a * df_b

# Two calls to to_pcollection with the same Dataframe should produce the

# same PCollection(s)

pc_2a_, pc_ab_ = convert.to_pcollection(df_2a, df_ab)

pc_3a, pc_2a, pc_ab = convert.to_pcollection(df_3a, df_2a, df_ab)

self.assertIs(pc_2a, pc_2a_)

self.assertIs(pc_ab, pc_ab_)

self.assertIsNot(pc_3a, pc_2a)

self.assertIsNot(pc_3a, pc_ab)

# The same conversions without the unbatching transform should also cache

# PCollections

pc_2a_pandas_, pc_ab_pandas_ = convert.to_pcollection(df_2a, df_ab,

yield_elements='pandas')

pc_3a_pandas, pc_2a_pandas, pc_ab_pandas = convert.to_pcollection(df_3a,

df_2a,

df_ab,

yield_elements='pandas')

self.assertIs(pc_2a_pandas, pc_2a_pandas_)

self.assertIs(pc_ab_pandas, pc_ab_pandas_)

self.assertIsNot(pc_3a_pandas, pc_2a_pandas)

self.assertIsNot(pc_3a_pandas, pc_ab_pandas)

# .. but the cached PCollections should be different

self.assertIsNot(pc_2a_pandas, pc_2a)

self.assertIsNot(pc_ab_pandas, pc_ab)

self.assertIsNot(pc_3a_pandas, pc_3a)

def test_convert_memoization_clears_cache(self):

# This test re-runs the other memoization test, and makes sure that the

# cache is cleaned up with the pipeline. Otherwise there would be concerns

# of it growing without bound.

import gc

# Make sure cache is clear

gc.collect()

self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 0)

# Disable GC so it doesn't run pre-emptively, confounding assertions about

# cache size

gc.disable()

# Also disable logging, as some implementations may artificially extend

# the life of objects.

import logging

logging.disable(logging.INFO)

try:

self.test_convert_memoization()

self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 3)

gc.collect()

# PCollections should be removed from cache after pipelines go out of

# scope and are GC'd

self.assertEqual(len(convert.TO_PCOLLECTION_CACHE), 0)

finally:

# Always re-enable GC and logging

gc.enable()

logging.disable(logging.NOTSET)

def test_auto_convert(self):

class MySchemaTransform(beam.PTransform):

def expand(self, pcoll):

return pcoll | beam.Map(

lambda x: beam.Row(

a=x.n**2 - x.m**2, b=2 * x.m * x.n, c=x.n**2 + x.m**2))

with beam.Pipeline() as p:

pc_mn = p | beam.Create([

(1, 2), (2, 3), (3, 10)

]) | beam.MapTuple(lambda m, n: beam.Row(m=m, n=n))

df_mn = convert.to_dataframe(pc_mn)

# Apply a transform directly to a dataframe to get another dataframe.

df_abc = df_mn | MySchemaTransform()

pc_abc = convert.to_pcollection(df_abc) | beam.Map(tuple)

assert_that(pc_abc, equal_to([(3, 4, 5), (5, 12, 13), (91, 60, 109)]))

if __name__ == '__main__':

unittest.main()