实验 15：分布式连接

from __future__ import (
    annotations,
)  # postponed evaluation of annotations https://peps.python.org/pep-0563/

import sys, pathlib

# Put the repo root (the directory containing `system/`) on sys.path so
# `from system...` works whether this notebook is opened from its own
# folder or from the repo root.
for _p in (pathlib.Path.cwd(), *pathlib.Path.cwd().parents):
    if (_p / "system").is_dir():
        if str(_p) not in sys.path:
            sys.path.insert(0, str(_p))
        break

import duckdb
import pandas as pd
from typing import Any
from faker import Faker
from system.data_classes import *

Faker.seed(42)
fake = Faker()

number_of_tuples = 10000

# create some fake data mimicking tuples in a relation
persons = [Person(fake.name(), fake.date()) for _ in range(number_of_tuples)]
books = [
    Book(
        fake.sentence(),
        persons[i % number_of_tuples].name,
        fake.pyfloat(min_value=0, max_value=100, right_digits=2),
    )
    for i in range(number_of_tuples)
]

# convert to Pandas DataFrame
persons_df = pd.DataFrame(persons)
books_df = pd.DataFrame(books)
display(persons_df[:5])
display(books_df[:5])

class Server:
    """A simple server that contains a duckdb connection.

    Attributes:
        con (duckdb.duckdb.DuckDBPyConnection): The duckdb connection.
        local_data (dict[str, pd.DataFrame): The local data.
    """

    def __init__(self, con: duckdb.duckdb.DuckDBPyConnection):
        self.con: duckdb.duckdb.DuckDBPyConnection = con  # DBMS running on this server
        self.local_data: dict[str, pd.DataFrame] = {}

    def import_data(self, tables: dict[str, pd.DataFrame]):
        """Import the data tables into the duckdb instance.

        :param tables: A dictionary of table name and data to import.
        """
        # duckdb resolves the literal `df_name` in the query below by the
        # Python variable name (replacement scan), so df_name is used even
        # though it never appears interpolated.
        # pylint: disable-next=unused-variable
        for table_name, df_name in tables.items():
            self.con.execute(f"CREATE TABLE {table_name} AS SELECT * FROM df_name")

    def execute_query(self, query: str, parameters: list[Any], name: str = "result"):
        """Execute the given `query` with its `parameters` on the duckdb instance. The result is stored in the local data area with the name `name` and registered in the duckdb instance.

        :param query: The query to execute.
        :param parameters: The parameters belonging to the query. Required to execute the query as a prepared statement.
        :param name: The name of the result table.
        """
        result_df = self.con.execute(query, parameters).fetchdf()  # prepared statement
        self.local_data[name] = result_df
        self.con.register(name, result_df)

    def get_local_data(self, name: str):
        """Retrieve the data table specified by the name `name` from the local data area.

        :param name: The name of the data table to retrieve.
        :return: The data table specified by the name `name`.
        """
        try:
            data_df = self.local_data[name]
            return data_df
        except KeyError:
            print(f"No local data for {name}.")
            return

    def receive(self, name: str, data_df: pd.DataFrame):
        """Receive a table with the name `name` and the data given by `data_df`. Simulates receiving data by registering the data in the duckdb instance.

        :param name: The name of the data table.
        :param data_df: The data to be received.
        """
        self.con.register(
            name, data_df
        )  # simulate receiving data by registering the data in the duckdb instance

    def send(self, other: Server, local_data_name: str):
        """Sends the data specified by `local_data_name` from the servers local data area to the `other` server.

        :param other: Server to send the data to.
        :param local_data_name: The name of the local data table to send.
        """
        data_df = self.get_local_data(local_data_name)
        other.receive(
            local_data_name, data_df
        )  # simulate sending the data to another server

# Utility function to compute the size of a database table
def table_size(table: pd.DataFrame):
    """Returns the size of a pandas DataFrame in kilobytes. This function assumes 4 byte for numeric/non-string values and counts the length of the string otherwise.

    :param table: pandas DataFrame representing a table.
    :return: size of the table in kilobytes.
    """
    size = 0  # byte
    for _, row in table.iterrows():
        for e in row:
            if type(e) is str:
                size += len(e)  # length of string
            else:
                size += 4  # assume 4 byte for numeric/non-string values
    return size / 1024

# Setup server 1
s1 = Server(duckdb.connect(database=":memory:"))
s1.import_data({"persons": persons_df})  # s1 contains entire Persons table

# Setup server 2
s2 = Server(duckdb.connect(database=":memory:"))
s2.import_data({"books": books_df})  # s2 contains entire Books table

# 1. Extract Books table on s2 and store it in the local data area of s2 with the name "books_local"
s2.execute_query("SELECT * FROM books", [], "books_local")

# 2. Send copy of the Books table to s1
s2.send(s1, "books_local")
print(f"Sending books table of size {table_size(s2.get_local_data("books_local"))} KB.")

# 3. s1 now contains the Books data in its local data area -> perform join between Persons and Books
s1.execute_query(
    """SELECT *
       FROM persons JOIN books_local on name = author
       WHERE birthday < ?
    """,
    ["2001-01-01"],
    "result",  # name used to store the result in the local data area of s1
)
result_entire_table = s1.get_local_data(
    "result"
)  # retrieve the previously computed result
print("Number of results:", result_entire_table.shape[0])

# 1. Extract the primary key of all persons born before the 21th century on s1 and store the result in the local data area of s1 with the name "persons_filtered"
s1.execute_query(
    "SELECT DISTINCT name FROM persons WHERE birthday < ?",
    ["2001-01-01"],
    "persons_filtered",
)

# 2. Send copy of filtered persons to S2
s1.send(s2, "persons_filtered")
persons_filtered_size = table_size(s1.get_local_data("persons_filtered"))
print(f"Sending filtered persons table of size {persons_filtered_size} KB.")

# 3. Extract books written by persons born before the 21th century
s2.execute_query(
    "SELECT books.* FROM books JOIN persons_filtered ON author = name",
    [],
    "books_filtered",
)  # semi-join

# 4. Send copy of filtered books to s1
s2.send(s1, "books_filtered")
books_filtered_size = table_size(s2.get_local_data("books_filtered"))
print(f"Sending filtered books table of size {books_filtered_size} KB.")

# 5. Execute final query, note that we could also reuse the pre-filtered persons table
s1.execute_query(
    """SELECT *
       FROM persons JOIN books_filtered on name = author
       WHERE persons.birthday < ?
    """,
    ["2001-01-01"],
    "result",
)
result_semi_join = s1.get_local_data(
    "result"
)  # retrieve the previously computed result
print("Number of results:", result_semi_join.shape[0])

def partition(data: list, attribute_name: str):
    """Partitions the input data based on the first character of the last name.
    :param data: List of entities to partition.
    :param attribute_name: Name of the attribute on which to partition the data.
    :return: Two pandas.DataFrames containing both partitions.
    """
    p1 = []
    p2 = []
    for e in data:
        last_name = getattr(e, attribute_name).split(" ")[1]
        if last_name[0] < "M":
            p1.append(e)
        else:
            p2.append(e)
    return pd.DataFrame(p1), pd.DataFrame(p2)

# Partition the data based on the join attribute
s1_persons_df, s2_persons_df = partition(persons, "name")
s1_books_df, s2_books_df = partition(books, "author")

# Setup server 1
s1 = Server(duckdb.connect(database=":memory:"))
s1.import_data(
    {"persons": s1_persons_df, "books": s1_books_df}
)  # s1 contains a partition of Books and Persons

# Setup server 2
s2 = Server(duckdb.connect(database=":memory:"))
s2.import_data(
    {"persons": s2_persons_df, "books": s2_books_df}
)  # s2 contains a partition of Books and Persons

# 1. Compute the partial join result locally on each node
s1.execute_query(
    "SELECT * FROM persons JOIN books ON name = author WHERE birthday < ?",
    ["2001-01-01"],
    "result_s1",
)
s2.execute_query(
    "SELECT * FROM persons JOIN books ON name = author WHERE birthday < ?",
    ["2001-01-01"],
    "result_s2",
)

# 2. Send `result_s1` to s2 since it is smaller than `result_s2`
s1.send(s2, "result_s1")
result_s1_size = table_size(s1.get_local_data("result_s1"))
print(f"Sending partial result computed on s1 of size {result_s1_size} KB to s2.")

# Compute final result by combining both partial results
s2.execute_query(
    "SELECT * FROM result_s1 UNION SELECT * FROM result_s2", [], "result_final"
)
result_local_joins = s2.get_local_data("result_final")
print("Number of results:", result_local_joins.shape[0])

import unittest
from pandas.testing import assert_frame_equal


class MyTest(unittest.TestCase):

    def test_result_sets(self):
        # Compare the three previously computed result sets.
        # Sort each table on all columns to ensure that the order of rows match
        result_11 = result_entire_table.sort_values(
            by=result_entire_table.columns.tolist()
        ).reset_index(
            drop=True
        )  # result computed by approach 1.1
        result_12 = result_semi_join.sort_values(
            by=result_semi_join.columns.tolist()
        ).reset_index(
            drop=True
        )  # result computed by approach 1.2
        result_21 = result_local_joins.sort_values(
            by=result_local_joins.columns.tolist()
        ).reset_index(
            drop=True
        )  # result computed by approach 2.1

        assert_frame_equal(result_11, result_12)
        assert_frame_equal(result_12, result_21)


# Run unit tests
unittest.main(argv=["ignored", "-v"], defaultTest="MyTest", verbosity=2, exit=False)