data_handler.py

import duckdb
from huggingface_hub import HfFileSystem

def setup_duckdb_hf_connection():
    """
    Registers the Hugging Face filesystem with DuckDB to enable
    'hf://' URI support for querying data.
    """
    try:
        hf_fs = HfFileSystem()
        duckdb.register_filesystem(hf_fs)
        print("Successfully registered HfFileSystem with DuckDB.")
        return True
    except Exception as e:
        print(f"Failed to register HfFileSystem: {e}")
        return False

def query_hf_data(query: str):
    """
    Executes a SQL query against Hugging Face data using DuckDB.
    """
    setup_duckdb_hf_connection()
    try:
        result = duckdb.sql(query).df()
        return result
    except Exception as e:
        print(f"Error executing query: {e}")
        return None

if __name__ == "__main__":
    # Example usage:
    # query = "SELECT * FROM 'hf://datasets/Brettapps/my-bucket/data.parquet' LIMIT 10"
    # df = query_hf_data(query)
    # print(df)
    pass