#

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

#

# pytype: skip-file

import os

import shutil

import sys

import tempfile

import unittest

import uuid

from collections.abc import Sequence

from typing import NamedTuple

from typing import Union

import numpy as np

from parameterized import parameterized

import apache_beam as beam

from apache_beam.io.filesystems import FileSystems

# pylint: disable=wrong-import-position, ungrouped-imports

try:

import tensorflow as tf

from tensorflow_transform.tf_metadata import dataset_metadata

from tensorflow_transform.tf_metadata import schema_utils

from apache_beam.ml.transforms import handlers

from apache_beam.ml.transforms import tft

from apache_beam.ml.transforms.tft import TFTOperation

from apache_beam.testing.util import assert_that

from apache_beam.testing.util import equal_to

except ImportError:

tft = None # type: ignore[assignment]

if not tft:

raise unittest.SkipTest('tensorflow_transform is not installed.')

class _AddOperation(TFTOperation):

def apply_transform(self, inputs, output_column_name, **kwargs):

return {output_column_name: inputs + 1}

class _MultiplyOperation(TFTOperation):

def apply_transform(self, inputs, output_column_name, **kwargs):

return {output_column_name: inputs * 10}

class IntType(NamedTuple):

x: int

class ListIntType(NamedTuple):

x: list[int]

class NumpyType(NamedTuple):

x: np.int64

class TFTProcessHandlerTest(unittest.TestCase):

def setUp(self) -> None:

self.artifact_location = tempfile.mkdtemp()

def tearDown(self):

shutil.rmtree(self.artifact_location)

@parameterized.expand([

({

'x': 1, 'y': 2

}, ['x'], {

'x': 20, 'y': 2

}),

({

'x': 1, 'y': 2

}, ['x', 'y'], {

'x': 20, 'y': 30

}),

])

def test_tft_operation_preprocessing_fn(

self, inputs, columns, expected_result):

add_fn = _AddOperation(columns=columns)

mul_fn = _MultiplyOperation(columns=columns)

process_handler = handlers.TFTProcessHandler(

transforms=[add_fn, mul_fn], artifact_location=self.artifact_location)

actual_result = process_handler.process_data_fn(inputs)

self.assertDictEqual(actual_result, expected_result)

def test_input_type_from_schema_named_tuple_pcoll(self):

data = [{'x': 1}]

with beam.Pipeline() as p:

data = (

p | beam.Create(data)

| beam.Map(lambda x: IntType(**x)).with_output_types(IntType))

element_type = data.element_type

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

inferred_input_type = process_handler._map_column_names_to_types(

element_type)

expected_input_type = dict(x=list[int])

self.assertEqual(inferred_input_type, expected_input_type)

def test_input_type_from_schema_named_tuple_pcoll_list(self):

data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}]

with beam.Pipeline() as p:

data = (

p | beam.Create(data)

| beam.Map(lambda x: ListIntType(**x)).with_output_types(ListIntType))

element_type = data.element_type

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

inferred_input_type = process_handler._map_column_names_to_types(

element_type)

expected_input_type = dict(x=list[int])

self.assertEqual(inferred_input_type, expected_input_type)

def test_input_type_from_row_type_pcoll(self):

data = [{'x': 1}]

with beam.Pipeline() as p:

data = (

p | beam.Create(data)

| beam.Map(lambda ele: beam.Row(x=int(ele['x']))))

element_type = data.element_type

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

inferred_input_type = process_handler._map_column_names_to_types(

element_type)

expected_input_type = dict(x=list[int])

self.assertEqual(inferred_input_type, expected_input_type)

def test_input_type_from_row_type_pcoll_list(self):

data = [{'x': [1, 2, 3]}, {'x': [4, 5, 6]}]

with beam.Pipeline() as p:

data = (

p | beam.Create(data)

| beam.Map(lambda ele: beam.Row(x=list(ele['x']))).with_output_types(

beam.row_type.RowTypeConstraint.from_fields([('x', list[int])])))

element_type = data.element_type

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

inferred_input_type = process_handler._map_column_names_to_types(

element_type)

expected_input_type = dict(x=list[int])

self.assertEqual(inferred_input_type, expected_input_type)

def test_input_type_from_named_tuple_pcoll_numpy(self):

np_data = [{

'x': np.array([1, 2, 3], dtype=np.int64)

}, {

'x': np.array([4, 5, 6], dtype=np.int64)

}]

with beam.Pipeline() as p:

data = (

p | beam.Create(np_data)

| beam.Map(lambda x: NumpyType(**x)).with_output_types(NumpyType))

element_type = data.element_type

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

inferred_input_type = process_handler._map_column_names_to_types(

element_type)

expected_type = dict(x=np.int64)

self.assertEqual(inferred_input_type, expected_type)

def test_tensorflow_raw_data_metadata_primitive_types(self):

input_types = dict(x=int, y=float, k=bytes, l=str)

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

for col_name, typ in input_types.items():

feature_spec = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

self.assertEqual(

handlers._default_type_to_tensor_type_map[typ], feature_spec.dtype)

self.assertIsInstance(feature_spec, tf.io.VarLenFeature)

def test_tensorflow_raw_data_metadata_primitive_types_in_containers(self):

input_types = dict([("x", list[int]), ("y", list[float]),

("k", list[bytes]), ("l", list[str])])

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

for col_name, typ in input_types.items():

feature_spec = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

self.assertIsInstance(feature_spec, tf.io.VarLenFeature)

@unittest.skipIf(sys.version_info < (3, 9), "not supported in python<3.9")

def test_tensorflow_raw_data_metadata_primitive_native_container_types(self):

input_types = dict([("x", list[int]), ("y", list[float]),

("k", list[bytes]), ("l", list[str])])

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

for col_name, typ in input_types.items():

feature_spec = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

self.assertIsInstance(feature_spec, tf.io.VarLenFeature)

def test_tensorflow_raw_data_metadata_numpy_types(self):

input_types = dict(x=np.int64, y=np.float32, z=list[np.int64])

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

for col_name, typ in input_types.items():

feature_spec = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

self.assertIsInstance(feature_spec, tf.io.VarLenFeature)

def test_tensorflow_raw_data_metadata_union_type_in_single_column(self):

input_types = dict(x=Union[int, float])

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

with self.assertRaises(TypeError):

for col_name, typ in input_types.items():

_ = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

def test_tensorflow_raw_data_metadata_dtypes(self):

input_types = dict(x=np.int32, y=np.float64)

expected_dtype = dict(x=np.int64, y=np.float32)

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

for col_name, typ in input_types.items():

feature_spec = process_handler._get_raw_data_feature_spec_per_column(

typ=typ, col_name=col_name)

self.assertEqual(expected_dtype[col_name], feature_spec.dtype)

def test_tft_process_handler_default_transform_types(self):

transforms = [

tft.ScaleTo01(columns=['x']),

tft.ScaleToZScore(columns=['y']),

tft.Bucketize(columns=['z'], num_buckets=2),

tft.ComputeAndApplyVocabulary(columns=['w'])

]

process_handler = handlers.TFTProcessHandler(

transforms=transforms, artifact_location=self.artifact_location)

column_type_mapping = (

process_handler._map_column_names_to_types_from_transforms())

expected_column_type_mapping = {

'x': float, 'y': float, 'z': float, 'w': str

}

self.assertDictEqual(column_type_mapping, expected_column_type_mapping)

expected_tft_raw_data_feature_spec = {

'x': tf.io.VarLenFeature(tf.float32),

'y': tf.io.VarLenFeature(tf.float32),

'z': tf.io.VarLenFeature(tf.float32),

'w': tf.io.VarLenFeature(tf.string)

}

actual_tft_raw_data_feature_spec = (

process_handler.get_raw_data_feature_spec(column_type_mapping))

self.assertDictEqual(

actual_tft_raw_data_feature_spec, expected_tft_raw_data_feature_spec)

def test_tft_process_handler_transformed_data_schema(self):

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location)

raw_data_feature_spec = {

'x': tf.io.VarLenFeature(tf.float32),

'y': tf.io.VarLenFeature(tf.float32),

'z': tf.io.VarLenFeature(tf.string),

}

raw_data_metadata = dataset_metadata.DatasetMetadata(

schema_utils.schema_from_feature_spec(raw_data_feature_spec))

expected_transformed_data_schema = {

'x': Sequence[np.float32],

'y': Sequence[np.float32],

'z': Sequence[bytes]

}

actual_transformed_data_schema = (

process_handler._get_transformed_data_schema(raw_data_metadata))

self.assertDictEqual(

actual_transformed_data_schema, expected_transformed_data_schema)

def test_tft_process_handler_verify_artifacts(self):

with beam.Pipeline() as p:

raw_data = (

p

| beam.Create([{

'x': np.array([1, 3])

}, {

'x': np.array([4, 6])

}]))

process_handler = handlers.TFTProcessHandler(

transforms=[tft.ScaleTo01(columns=['x'])],

artifact_location=self.artifact_location,

)

_ = raw_data | process_handler

self.assertTrue(

FileSystems.exists(

# To check the gcs directory with FileSystems, the dir path must

# end with /

os.path.join(

self.artifact_location,

handlers.RAW_DATA_METADATA_DIR + '/')))

self.assertTrue(

FileSystems.exists(

os.path.join(

self.artifact_location,

handlers.RAW_DATA_METADATA_DIR,

handlers.SCHEMA_FILE)))

with beam.Pipeline() as p:

raw_data = (p | beam.Create([{'x': np.array([2, 5])}]))

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location, artifact_mode='consume')

transformed_data = raw_data | process_handler

transformed_data |= beam.Map(lambda x: x.x)

# the previous min is 1 and max is 6. So this should scale by (1, 6)

assert_that(

transformed_data,

equal_to([np.array([0.2, 0.8], dtype=np.float32)],

equals_fn=np.array_equal))

def test_tft_process_handler_unused_column(self):

data = [

{

'x': [5, 1], 'y': 2

},

{

'x': [8, 4], 'y': 5

},

{

'x': [9, 5], 'y': 6

},

{

'x': [10, 6], 'y': 7

},

{

'x': [11, 7], 'y': 8

},

{

'x': [12, 8], 'y': 9

},

{

'x': [13, 9], 'y': 10

},

{

'x': [14, 10], 'y': 11

},

{

'x': [15, 11], 'y': 12

},

{

'x': [16, 12], 'y': 13

},

{

'x': [17, 13], 'y': 14

},

{

'x': [18, 14], 'y': 15

},

{

'x': [19, 15], 'y': 16

},

{

'x': [20, 16], 'y': 17

},

{

'x': [21, 17], 'y': 18

},

{

'x': [22, 18], 'y': 19

},

{

'x': [23, 19], 'y': 20

},

{

'x': [24, 20], 'y': 21

},

{

'x': [25, 21], 'y': 22

},

]

# pylint: disable=line-too-long

expected_data = [

beam.Row(

x=np.array([0.16666667, 0.], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=2),

beam.Row(

x=np.array([0.29166666, 0.125], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=5),

beam.Row(

x=np.array([0.33333334, 0.16666667], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=6),

beam.Row(

x=np.array([0.375, 0.20833333], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=7),

beam.Row(

x=np.array([0.41666666, 0.25], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=8),

beam.Row(

x=np.array([0.45833334, 0.29166666], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=9),

beam.Row(

x=np.array([0.5, 0.33333334], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=10),

beam.Row(

x=np.array([0.5416667, 0.375], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=11),

beam.Row(

x=np.array([0.5833333, 0.41666666], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=12),

beam.Row(

x=np.array([0.625, 0.45833334], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=13),

beam.Row(

x=np.array([0.6666667, 0.5], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=14),

beam.Row(

x=np.array([0.7083333, 0.5416667], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=15),

beam.Row(

x=np.array([0.75, 0.5833333], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=16),

beam.Row(

x=np.array([0.7916667, 0.625], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=17),

beam.Row(

x=np.array([0.8333333, 0.6666667], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=18),

beam.Row(

x=np.array([0.875, 0.7083333], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=19),

beam.Row(

x=np.array([0.9166667, 0.75], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=20),

beam.Row(

x=np.array([0.9583333, 0.7916667], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=21),

beam.Row(

x=np.array([1., 0.8333333], dtype=np.float32),

x_max=np.array([25.], dtype=np.float32),

x_min=np.array([1.], dtype=np.float32),

y=22),

]

# pylint: enable=line-too-long

expected_data_x = [row.x for row in expected_data]

expected_data_y = [row.y for row in expected_data]

scale_to_0_1_fn = tft.ScaleTo01(columns=['x'])

with beam.Pipeline() as p:

raw_data = p | beam.Create(data)

process_handler = handlers.TFTProcessHandler(

transforms=[scale_to_0_1_fn],

artifact_location=self.artifact_location,

)

transformed_pcoll = raw_data | process_handler

transformed_pcoll_x = transformed_pcoll | beam.Map(lambda x: x.x)

transformed_pcoll_y = transformed_pcoll | beam.Map(lambda x: x.y)

assert_that(

transformed_pcoll_x,

equal_to(expected_data_x, equals_fn=np.array_equal),

label='transformed data')

assert_that(

transformed_pcoll_y,

equal_to(expected_data_y, equals_fn=np.array_equal),

label='unused column')

def test_consume_mode_with_extra_columns_in_the_input(self):

with beam.Pipeline() as p:

raw_data = (

p

| beam.Create([{

'x': np.array([1, 3])

}, {

'x': np.array([4, 6])

}]))

process_handler = handlers.TFTProcessHandler(

transforms=[tft.ScaleTo01(columns=['x'])],

artifact_location=self.artifact_location,

)

_ = raw_data | process_handler

test_data = [{

'x': np.array([2, 5]), 'y': np.array([1, 2]), 'z': 'fake_string'

},

{

'x': np.array([1, 10]),

'y': np.array([2, 3]),

'z': 'fake_string2'

}]

expected_test_data = [{

'x': np.array([0.2, 0.8], dtype=np.float32),

'y': np.array([1, 2]),

'z': 'fake_string'

},

{

'x': np.array([0., 1.8], dtype=np.float32),

'y': np.array([2, 3]),

'z': 'fake_string2'

}]

expected_test_data_x = [row['x'] for row in expected_test_data]

expected_test_data_y = [row['y'] for row in expected_test_data]

expected_test_data_z = [row['z'] for row in expected_test_data]

with beam.Pipeline() as p:

raw_data = (p | beam.Create(test_data))

process_handler = handlers.TFTProcessHandler(

artifact_location=self.artifact_location, artifact_mode='consume')

transformed_data = raw_data | process_handler

transformed_data_x = transformed_data | beam.Map(lambda x: x.x)

transformed_data_y = transformed_data | beam.Map(lambda x: x.y)

transformed_data_z = transformed_data | beam.Map(lambda x: x.z)

assert_that(

transformed_data_x,

equal_to(expected_test_data_x, equals_fn=np.array_equal),

label='transformed data')

assert_that(

transformed_data_y,

equal_to(expected_test_data_y, equals_fn=np.array_equal),

label='unused column: y')

assert_that(

transformed_data_z,

equal_to(expected_test_data_z, equals_fn=np.array_equal),

label='unused column: z')

def test_handler_with_same_input_elements(self):

with beam.Pipeline() as p:

data = [

{

'x': 'I'

},

{

'x': 'love'

},

{

'x': 'Beam'

},

{

'x': 'Beam'

},

{

'x': 'is'

},

{

'x': 'awesome'

},

]

raw_data = (p | beam.Create(data))

process_handler = handlers.TFTProcessHandler(

transforms=[tft.ComputeAndApplyVocabulary(columns=['x'])],

artifact_location=self.artifact_location,

)

transformed_data = raw_data | process_handler

expected_data = [

beam.Row(x=np.array([4])),

beam.Row(x=np.array([1])),

beam.Row(x=np.array([0])),

beam.Row(x=np.array([0])),

beam.Row(x=np.array([2])),

beam.Row(x=np.array([3])),

]

expected_data_x = [row.x for row in expected_data]

actual_data_x = transformed_data | beam.Map(lambda x: x.x)

assert_that(

actual_data_x,

equal_to(expected_data_x, equals_fn=np.array_equal),

label='transformed data')

class TFTProcessHandlerTestWithGCSLocation(TFTProcessHandlerTest):

def setUp(self) -> None:

self.artifact_location = self.gcs_artifact_location = os.path.join(

'gs://temp-storage-for-perf-tests/tft_handler', uuid.uuid4().hex)

def tearDown(self):

pass

if __name__ == '__main__':

unittest.main()