code/collaborative-optimizer/examples/user-defined-functions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c289b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment Import\n",
    "from experiment_graph.execution_environment import ExecutionEnvironment\n",
    "from experiment_graph.graph.operations import UserDefinedFunction\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5da15db1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ClipDataset(UserDefinedFunction):\n",
    "    def __init__(self, lower=None, upper=None, axis=None):\n",
    "        super().__init__(return_type='Dataset')\n",
    "        self.lower = lower\n",
    "        self.upper= upper\n",
    "        self.axis = axis\n",
    "        \n",
    "    def run(self, underlying_data):\n",
    "        # here the underlying_data is a pandas dataframe and we are directly calling the pandas clip function\n",
    "        return underlying_data.clip(lower=self.lower, upper=self.upper, axis=self.axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ca669fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomProject(UserDefinedFunction):\n",
    "    def __init__(self, column):\n",
    "        super().__init__(return_type='Feature')\n",
    "        self.column = column\n",
    "        \n",
    "    def run(self, underlying_data):\n",
    "        # here the underlying_data is a pandas dataframe but the udf is expected to return a Feature (i.e., dataseries)\n",
    "        return underlying_data[self.column]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e02c98d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MedianFeature(UserDefinedFunction):\n",
    "    def __init__(self):\n",
    "        super().__init__(return_type='Agg')\n",
    "        \n",
    "    def run(self, underlying_data):\n",
    "        # here the underlying_data is pandas data series\n",
    "        return underlying_data.median()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5469d33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "execution_environment = ExecutionEnvironment()\n",
    "data = execution_environment.load('./data/sample.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "953cd7f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "clip_oper = ClipDataset(lower=2,upper=4)\n",
    "clipped = data.run_udf(clip_oper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdefa2ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "clipped.data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "108c2212",
   "metadata": {},
   "outputs": [],
   "source": [
    "projection = CustomProject('a')\n",
    "feature_a = clipped.run_udf(projection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93444ee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_a.data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15395eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "median_feature=MedianFeature()\n",
    "median_value = feature_a.run_udf(median_feature)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "493dadd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "median_value.data()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}