What does the Spark Expectations Plugin Ecosystem Look Like

[newfront]

What are the requirements needed in order to create a good plugin architecture?

[newfront]

with @asingamaneni @aditya-chaturvedi @phanikumarvemuri

[asingamaneni]

Well in general, for a good-plugin architecture below are design, implementation and code structure considerations which we could take.
@newfront @aditya-chaturvedi @Umeshsp22 @phanikumarvemuri , please add more if I miss anything below

Design Considerations

1. Modularity: The core application should be separated from the plugins, so that each evolve independently. Separate folder structure should do this for us.
2. Loose Coupling: The core application should not be tightly integrated with any particular plugin, but a default plugin loading could be provided if none is provided at run time.
3. Extensibility: It should be easy to add or remove plugins without modifying the core application.
4. Consistency: A standard interface should be defined for each functionality, that all plugins must adhere to.
5. Version Compatability: The architecture should be designed to handle different versions of plugins, this could be tricky

Implementation Considerations

1. Dynamic Loading: Use Python's importlib to dynamically load plugins at runtime
2. Configuration: Could use a configuration file to manage which plugins are enabled or disabled
3. Error Handling: Robust error handling to manage failures in plugins.
4. Logging: Implement logging within the plugin architecture within in the plugin architecture to track performance and errors.
5. Security: Ensure that plugins can't execute malicious code or access sentive parts of the core application.
6. Documentation: Detailed documentation for plugin developers explaining how to create a new plugin, what interfaces to implement etc.

Code Structure

1. Plugin Interface: Define an abstract base class that all plugins must inherit from

from abc import ABC, abstractmethod

class PluginBase(ABC):
	@abstractmethod
	def execute(self, data):
		pass

2. Plugin Manager: A class to manage the loading, unloading and execution of plugins

import importlib

class PluginManager:
	def __init__(self):
		self.plugins = {}

	def load_plugins(self, plugin_name):
		module = importlib.import_module(plugin_name)
		self.plugins[plugin_name] = module.Plugin()

	def unload_plugin(self, plugin_name):
		del self.plugins[plugin_name]

	def execute_plugins(self, data):
		for plugin in self.plugins.values():
			plugin.execute(data)

3. Plugin Implementation: Each plugin will implement the PluginBase class

class MyPlugin(PluginBase):
	def execute(self, data):
		print(f"plugin executed with data: {data}")

Example implementation for our usecase in Spark-Expectations could be as below:

Define a Standard Reader Interface

First, define an abstract base class that all reader plugins must implement. This ensures that they all have the same interface.

from abc import ABC, abstractmethod
class ReaderPlugin(ABC):
    @abstractmethod
    def read(self, *args, **kwargs):
        pass

Implement the Default Reader

Default reader for reading from Delta can be a concrete implementation of this interface.

class DeltaReader(ReaderPlugin):
    def read(self, path):
        # Delta reading logic here

Create a Reader Manager

Create a manager class that will dynamically load the appropriate reader plugin at runtime.

import importlib

class ReaderManager:
    def __init__(self):
        self.readers = {'delta': DeltaReader()}

    def register_reader(self, name, reader_class_path):
        module_name, class_name = reader_class_path.rsplit('.', 1)
        module = importlib.import_module(module_name)
        reader_class = getattr(module, class_name)

        if not issubclass(reader_class, ReaderPlugin):
            raise ValueError("Invalid reader plugin")

        self.readers[name] = reader_class()

    def get_reader(self, name):
        return self.readers.get(name)

User-Defined Reader Plugin

Users can create their own reader plugins by implementing the ReaderPlugin interface. For example, a Snowflake reader might look like this:

class SnowflakeReader(ReaderPlugin):
    def read(self, query, connection_params):
        # Snowflake reading logic here

Registering User-Defined Plugins

Users can register their custom reader plugins with ReaderManager.

reader_manager = ReaderManager()
reader_manager.register_reader('snowflake', 'my_package.SnowflakeReader')

Using the Reader at Runtime

When spark-expectations library needs to read data, it can now use the ReaderManager to get the appropriate reader and use it.

def read_data(source_type, *args, **kwargs):
    reader = reader_manager.get_reader(source_type)
    if reader:
        return reader.read(*args, **kwargs)
    else:
        raise ValueError(f"No reader registered for {source_type}")

This way, our library remains open for extension but closed for modification, adhering to the Open/Closed Principle. Users can easily extend its functionality to read from different data sources without having to modify its core logic.

[newfront]

I like the use of the ReaderManager for loading up additional readers. Do you see use cases where there will be the need to connect to multiple backend data sources for doing multi-shot expectations management? I could see this in the case where people want to check that the results of an import from SourceA -> SourceB worked as expected.

df1 = (
  reader_manager
    .for(sources.SNOWFLAKE, 'my_package.SnowflakeReader',opts)
    .read(table|None)
    .options(sf_read_opts|None)
    .execute())

df2 = (
  reader_manager
    .for(sources.DELTA, 'delta_common_dbxA', opts)
    .read(table|None)
    .options(uc_read_opts|None)
    .execute())

# now run expectations

Leaning on the DataSource api : (DataFrameReader, DataStreamReader, org.apache.spark.sql.connector.catalog.{Table}, etc)

If we are leaning into the spark ecosystem anyways, then a reader can be based on the DataSource api. So the expectation that spark_session.read is already a given. So as long as the plugin just does some decoration over the SparkSession.read method then complexity will be low.

Additional thought: If the ReaderManager is initialized with the active SparkSession, then each individual reader can use a new_session so the internal Spark RuntimeConf doesn't collide.

base_spark_session: SparkSession = (SparkSession....getOrCreate() | SparkSession.getActiveSession...)

reader_manager = ReaderManager(spark_session=base_spark_session)

note: Registering a plugin that requires a separate spark.sql.catalog or independent Spark RuntimeConf

# https://github.com/snowflakedb/spark-snowflake/tree/v2.12.0-spark_3.4
# spark-shell --packages net.snowflake:snowflake-jdbc:3.13.22,net.snowflake:spark-snowflake_2.12:2.12.0-spark_3.4
# https://docs.snowflake.com/en/user-guide/spark-connector-use#setting-configuration-options-for-the-connector
(reader_manager
  .register('snowflake', 'conn1', config={#default_options}, new_spark_session=True) 
  })

sf_reader: WrappedDataFrameReader = (
  reader_manager
    .for("snowflake")
    .get_reader('conn1')
    .withOptions(**{'sfUrl': 'warehouse_url', 'sfUser': 'sfAuthenticator', 'sfToken': secrets.get('token_sf_tests'), 'sfDatabase': 'db', 'sfSchema': '...'))

sf_df: DataFrame: sf_reader.read().execute()

^^ Not sure if it would be helpful to wrap the DataFrameReader or provide direct access. There are some benefits to being able to provide an intermediate builder. Just spitballing here.

Test that all packages and (in case of SF) jdbc jars are on the path for the version...

ReaderManager.for("snowflake").get_reader('conn1').test() # raises exception if connection can't be made

[asingamaneni]

I like the idea of - WrappedDataFrameReader and providing source in there that could lead to the choosing of the implementation.
For df1 and df2, I am not sure - if we need to extend the implementation from spark-expectations, because the users will use native spark, to load the dataframes that are needed for the spark-expectations rules to run.

[phanikumarvemuri]

I completely agree with design considerations, but for implementation I was thinking more a Service Loader kind of approach so that the users don't need to register the plugins but in-fact the core module can load the class that provide implementation dynamically.
Similar to this
Now in order to load the class, the user has to pass the mapping between source and class name or we can load all the classes that extend a specific abstract classs like "all the classes that extend Reader class".

import importlib
import pkgutil
import inspect
from abc import ABC, abstractmethod


class SparkExpectationsSource(ABC):

    @abstractmethod
    def read(self, *args, **kwargs):
        pass

    @abstractmethod
    def write(self, *args, **kwargs):
        pass

class ServiceLoader:
    def __init__(self, base_class):
        self.base_class = base_class
        self.providers = self._load_all_providers()

    def _load_all_providers(self):
        providers = {}
        for importer, modname, ispkg in pkgutil.iter_modules():
            try:
                module = importer.find_module(modname).load_module(modname)
                for name, obj in inspect.getmembers(module):
                    if (
                            inspect.isclass(obj)
                            and issubclass(obj, self.base_class)
                            and obj is not self.base_class
                    ):
                        providers[obj.type] = obj()
            except ImportError:
                pass
        return providers

    def get_providers(self):
        return self.providers


#user needed type can be obtained from the config 
provider = ServiceLoader(SparkExpectationsSource).get_providers("user needed type")

provider.read()
provider.write()

[phanikumarvemuri]

The default implementation can come from the input dataframe itself ( check the lineage and fetch the source) assuming reading and writing will happen to same system.

[holdenk]

I think Security: Ensure that plugins can't execute malicious code or access sentive parts of the core application. might not be a reasonable goal/requirement. I think it's up to the user to only install trusted plugins.

[asingamaneni]

Makes Sense!

[newfront]

Yes and yes @holdenk. It also is nice to not be on the hook for security 😀

[asingamaneni]

I Liked both the approaches proposed by @newfront and @phanikumarvemuri , so tried to club them together in the original pattern. Let me know your thoughts

## import required Modules
from abc import ABC, abstractmethod
from pyspark.sql import SparkSession
import importlib
import pkgutil
import inspect


## Define the Abstract Reader Class (ReaderPlugin)
class ReaderPlugin(ABC):
    def __init__(self, spark_session, config):
        self.spark_session = spark_session
        self.config = config

    @abstractmethod
    def read(self, *args, **kwargs):
        pass


## Implement Delta Reader
class DeltaReader(ReaderPlugin):
    def read(self, path):
        print(f"Reading data from Delta at {path}")


## As per Phani's suggestion, modified a bit - OptimizedServiceLoader class
## We only scan required packages from a known class path rather than all the modules
## We lazily load the required plugin using get_provider function
## We store the providers which are already read
class ServiceLoader:
    def __init__(self, base_class, scan_packages=[]):
        self.base_class = base_class
        self.scan_packages = scan_packages
        self.providers = {}

    def _load_providers(self, package):
        for importer, modname, ispkg in pkgutil.iter_modules(package.__path__):
            try:
                module = importlib.import_module(f"{package.__name__}.{modname}")
                for name, obj in inspect.getmembers(module):
                    if inspect.isclass(obj) and issubclass(obj, self.base_class):
                        self.providers[obj.__name__] = obj
            except ImportError:
                pass

    def get_provider(self, name):
        if name not in self.providers:
            for package in self.scan_packages:
                self._load_providers(package)
        return self.providers.get(name)


## As per Scott's suggestion, implementing new spark_session per reader
class ReaderManager:
    def __init__(self, base_spark_session=None, scan_packages=[]):
        self.base_spark_session = base_spark_session or SparkSession.builder.getOrCreate()
        self.loader = ServiceLoader(ReaderPlugin, scan_packages=scan_packages)
        self.readers = {}

    def get_reader(self, name, config={}):
        if name not in self.readers:
            reader_class = self.loader.get_provider(name)
            if reader_class:
                new_session = self.base_spark_session.newSession() # Creates a new session
                self.readers[name] = reader_class(new_session, config) # Instantiate the reader with the new session
        return self.readers.get(name) # Return the cached reader if it exists, otherwise return None

User code could look like this with the above implementation

@se.with_expectations(
    reader_manager.get_reader('ReaderType', reader_config).get_rules_from_table(
        product_rules_table="pilot_nonpub.dq.dq_rules",
        table_name="pilot_nonpub.dq_employee.employee",
        dq_stats_table_name="pilot_nonpub.dq.dq_stats"
    ),
    # ... (rest of the existing code)
)
def build_new() -> DataFrame:
    # ... (rest of the existing code)

Sample Implementation for Snowflake

class SnowflakeReader(ReaderPlugin):
    def __init__(self, spark_session, config):
        super().__init__(spark_session, config)

    def read(self, query: str) -> DataFrame:
        url = self.config.get('url')
        user = self.config.get('user')
        password = self.config.get('password')
        database = self.config.get('database')
        schema = self.config.get('schema')
        warehouse = self.config.get('warehouse')

        return self.spark_session.read \
            .format("snowflake") \
            .option("sfUrl", url) \
            .option("sfDatabase", database) \
            .option("sfWarehouse", warehouse) \
            .option("sfSchema", schema) \
            .option("sfRole", "PUBLIC") \
            .option("sfFetchSize", 1000) \
            .option("user", user) \
            .option("password", password) \
            .option("dbtable", query) \
            .load()

User Code for snowflake

# Initialize ReaderManager
reader_manager = ReaderManager(base_spark_session=spark, scan_packages=['package_containing_SnowflakeReader'])

# Snowflake configuration
snowflake_config = {
    'url': 'snowflake_url',
    'user': 'username',
    'password': 'password',
    'database': 'database',
    'schema': 'schema',
    'warehouse': 'warehouse'
}

@se.with_expectations(
    reader_manager.get_reader('SnowflakeReader', snowflake_config).get_rules_from_table(
        product_rules_table="pilot_nonpub.dq.dq_rules",
        table_name="pilot_nonpub.dq_employee.employee",
        dq_stats_table_name="pilot_nonpub.dq.dq_stats"
    ),
    # ... (rest of the existing code)
)
def build_new() -> DataFrame:
    # ... (rest of the existing code)

[asingamaneni]

Me and @phanikumarvemuri met today and discussed the implementation. We re-thought the usage of plugins and may be plugins are not needed at all if we use native spark DataFrameWriter, DataFrameWriterV2

Decisions discussed:

1. User will provide the rules_df or rules_df_view as part of the SparkExpectations class initialization or defining with_expectations to override class object
2. We will provide writer_manager to the user through which they can provide all the spark writer options for target_table, error_table and stats_table and create separate objects which they can pass to with_expectatios. Using which we will call the DataFrameWriter.saveAsTable() with the relevant options provided to writer_manager