584. Find Customer Referee

Problem¶

Return the names of customers whose referee_id is either not 2 or NULL.

Input:

# Customer
---+------+-----------
id | name | referee_id
---+------+-----------
1  | Will | NULL
2  | Jane | NULL
3  | Alex | 2
5  | Zack | 1

Output:

----
name
----
Will
Jane
Zack

Key trick¶

Use an explicit NULL check:

referee_id <> 2 OR referee_id IS NULL

Trap¶

referee_id <> 2 alone does not match NULL in SQL.
Do not exclude the customer with id = 2; exclude customers whose referee_id = 2.

Why is it interesting?¶

This is a small test of SQL NULL semantics and three-valued logic.

SQL solution¶

-- Keep customers not referred by customer 2.
-- In SQL, NULL needs an explicit IS NULL check.
SELECT name
FROM Customer
WHERE referee_id <> 2
   OR referee_id IS NULL;

Pandas solution¶

import pandas as pd

def find_customer_referee(customer: pd.DataFrame) -> pd.DataFrame:
    # Keep rows where referee_id is not 2, including missing referee_id.
    mask = (customer["referee_id"] != 2) | customer["referee_id"].isna()

    # Return only the required column.
    return customer.loc[mask, ["name"]]

Pytest test¶

import sqlite3

import pandas as pd
import pytest


SQL_QUERY = """
SELECT name
FROM Customer
WHERE referee_id <> 2
   OR referee_id IS NULL;
"""


def find_customer_referee(customer: pd.DataFrame) -> pd.DataFrame:
    mask = (customer["referee_id"] != 2) | customer["referee_id"].isna()
    return customer.loc[mask, ["name"]]


@pytest.mark.parametrize(
    "rows, expected_names",
    [
        (
            [
                (1, "Will", None),
                (2, "Jane", None),
                (3, "Alex", 2),
                (4, "Bill", None),
                (5, "Zack", 1),
                (6, "Mark", 2),
            ],
            ["Bill", "Jane", "Will", "Zack"],
        ),
        (
            [
                (1, "A", None),
                (2, "B", None),
            ],
            ["A", "B"],
        ),
        (
            [
                (1, "A", 2),
                (2, "B", 2),
            ],
            [],
        ),
        (
            [
                (2, "CustomerTwo", 3),
                (3, "ReferredByTwo", 2),
                (4, "ReferredByOne", 1),
            ],
            ["CustomerTwo", "ReferredByOne"],
        ),
        (
            [],
            [],
        ),
    ],
)
def test_find_customer_referee_sql_and_pandas(rows, expected_names):
    conn = sqlite3.connect(":memory:")
    conn.execute(
        """
        CREATE TABLE Customer (
            id INTEGER,
            name TEXT,
            referee_id INTEGER
        );
        """
    )
    conn.executemany(
        """
        INSERT INTO Customer (id, name, referee_id)
        VALUES (?, ?, ?);
        """,
        rows,
    )

    sql_result = pd.read_sql_query(SQL_QUERY, conn)["name"].tolist()

    customer = pd.DataFrame(rows, columns=["id", "name", "referee_id"]).astype(
        {
            "id": "Int64",
            "name": "object",
            "referee_id": "Int64",
        }
    )
    pandas_result = find_customer_referee(customer)["name"].tolist()

    assert sorted(sql_result) == sorted(expected_names)
    assert sorted(pandas_result) == sorted(expected_names)

Comment on my solution¶

Your SQL solution is correct and handles NULL explicitly.
Your Pandas solution is also correct.
The important part is the isna() branch, because it mirrors SQL's IS NULL condition.

SELECT name
FROM Customer
WHERE referee_id <> 2
   OR referee_id IS NULL;

import pandas as pd

def find_customer_referee(customer: pd.DataFrame) -> pd.DataFrame:
    mask = (customer["referee_id"] != 2) | customer["referee_id"].isna()
    return customer.loc[mask, ["name"]]