ENH: Add support for standard sql datatypes

Author: zbrookle

dataframe_sql/grammar/sql.grammar

@@ -63,7 +63,22 @@ column_name: [NAME "."] NAME
 
 
 SELECT_CONSTRAINT.9: "ALL"i | "DISTINCT"i
-TYPENAME: "object"i | "int64"i | "float64"i | "bool"i | "datetime64"i | "timedelta[ns]"i | "category"i
+TYPENAME:  "object"i
+         | "varchar"i
+         | "int16"i
+         | "smallint"i
+         | "int32"i
+         | "int64"i
+         | "int"i
+         | "bigint"i
+         | "float16"i
+         | "float32"i
+         | "float64"i
+         | "float"i
+         | "bool"i
+         | "datetime64"i
+         | "timestamp"i
+         | "category"i
 ?aggregation: NAME -> aggregation_name
 alias: NAME -> alias_string
 _window_name: NAME

dataframe_sql/parsing/sql_parser.py

@@ -42,12 +42,34 @@
 )
 PANDAS_TYPE_PYTHON_TYPE_FUNCTION = {
     "object": str,
+    "string": str,
+    "int16": int,
+    "int32": int,
     "int64": int,
+    "float16": float,
+    "float32": float,
     "float64": float,
     "bool": bool,
 }
+
+TYPE_TO_PANDAS_TYPE = {
+    "varchar": "string",
+    "smallint": "int16",
+    "int": "int32",
+    "bigint": "int64",
+    "float": "float64",
+    "timestamp": "datetime64",
+    "datetime64": "datetime64",
+    "timedelta[ns]": "timedelta[ns]",
+    "category": "category",
+}
+
+for TYPE in PANDAS_TYPE_PYTHON_TYPE_FUNCTION:
+    TYPE_TO_PANDAS_TYPE[TYPE] = TYPE
+
 PANDAS_TYPE_TO_SQL_TYPE = {
     "object": String,
+    "string": String,
     "int64": Number,
     "float64": Number,
     "bool": Bool,
@@ -748,8 +770,6 @@ def rank(self, tokens, rank_function):
         :param rank_function: Function to be used in rank evaluation
         :return:
         """
-        print("yes")
-
         expressions = tokens[0]
         series_list = []
         order_list = []
@@ -846,7 +866,7 @@ def as_type(self, column_and_type):
         """
         column = column_and_type[0]
         typename = column_and_type[1]
-        column.typename = typename.value
+        column.typename = TYPE_TO_PANDAS_TYPE[typename.value]
         return column
 
     def literal_cast(self, value_and_type: list):
@@ -1218,7 +1238,10 @@ def handle_non_token_non_tree(query_info: QueryInfo, token, token_pos):
         if isinstance(token, Column):
             query_info.columns.append(token)
             query_info.column_selected[token.name] = True
+            # TODO Get rid of collecting this alias information since its part of the
+            #  column object
             if token.alias:
+                print(query_info.aliases)
                 query_info.aliases[token.name] = token.alias
 
         if isinstance(token, Expression):
@@ -1398,6 +1421,7 @@ def handle_columns(
         :param internal_transformer: Transformer to transform the where clauses
         :return:
         """
+        print(columns)
         where_value = None
         where_plan = ":"
         if where_expr is not None:
@@ -1415,6 +1439,7 @@ def handle_columns(
                 new_frame = first_frame.copy()
         else:
             column_names = []
+            final_names = []
             for column in columns:
                 true_column_name = self.column_name_map[column.table][
                     column.name.lower()
@@ -1426,13 +1451,20 @@ def handle_columns(
                 ):
                     aliases[true_column_name] = column.name
 
+                if column.alias:
+                    final_names.append(column.alias)
+                else:
+                    final_names.append(column.name)
+
+            print(final_names)
             if where_value is not None:
                 new_frame = first_frame.loc[where_value, column_names]
             else:
                 new_frame = first_frame.loc[:, column_names]
             execution_plan += f".loc[{where_plan}, {column_names}]"
             if aliases:
-                new_frame = new_frame.rename(columns=aliases)
+                # new_frame = new_frame.rename(columns=aliases)
+                new_frame.columns = final_names
                 execution_plan += f".rename(columns={aliases})"
 
         return new_frame, execution_plan

dataframe_sql/tests/pandas_sql_functionality_test.py

@@ -21,6 +21,7 @@
 from dataframe_sql.sql_objects import AmbiguousColumn
 from dataframe_sql.sql_select_query import TableInfo
 from dataframe_sql.tests.utils import (
+    AVOCADO,
     DIGIMON_MON_LIST,
     DIGIMON_MOVE_LIST,
     FOREST_FIRES,
@@ -1320,9 +1321,82 @@ def test_case_statement_with_same_conditions():
     tm.assert_frame_equal(pandas_frame, my_frame)
 
 
+@assert_state_not_change
+def test_multiple_aliases_same_column():
+    """
+    Test multiple aliases on the same column
+    :return:
+    """
+    my_frame = query(
+        """
+        select wind as my_wind, wind as also_the_wind, wind as yes_wind
+        from
+        forest_fires
+        """
+    )
+
+    pandas_frame = FOREST_FIRES[["wind"]].copy()
+    pandas_frame.loc[:, "my_wind"] = FOREST_FIRES["wind"].copy()
+    pandas_frame.loc[:, "also_the_wind"] = FOREST_FIRES["wind"]
+    pandas_frame.loc[:, "yes_wind"] = FOREST_FIRES["wind"]
+    pandas_frame = pandas_frame.drop(columns=["wind"])
+    tm.assert_frame_equal(pandas_frame, my_frame)
+
+
+@assert_state_not_change
+def test_sql_data_types():
+    """
+    Tests sql data types
+    :return:
+    """
+    my_frame = query(
+        """
+        select 
+            cast(avocado_id as object) as avocado_id_object,
+            cast(avocado_id as int16) as avocado_id_int16,
+            cast(avocado_id as smallint) as avocado_id_smallint,
+            cast(avocado_id as int32) as avocado_id_int32,
+            cast(avocado_id as int) as avocado_id_int,
+            cast(avocado_id as int64) as avocado_id_int64,
+            cast(avocado_id as bigint) as avocado_id_bigint,
+            cast(avocado_id as float) as avocado_id_float,
+            cast(avocado_id as float16) as avocado_id_float16,
+            cast(avocado_id as float32) as avocado_id_float32,
+            cast(avocado_id as float64) as avocado_id_float64,
+            cast(avocado_id as bool) as avocado_id_bool,
+            cast(avocado_id as category) as avocado_id_category,
+            cast(date as datetime64) as date,
+            cast(date as timestamp) as time,
+            cast(region as varchar) as region_varchar
+        from avocado
+        """
+    )
+
+    pandas_frame = AVOCADO.copy()[["avocado_id", "Date", "region"]]
+    pandas_frame["avocado_id_object"] = pandas_frame["avocado_id"].astype("object")
+    pandas_frame["avocado_id_int16"] = pandas_frame["avocado_id"].astype("int16")
+    pandas_frame["avocado_id_smallint"] = pandas_frame["avocado_id"].astype("int16")
+    pandas_frame["avocado_id_int32"] = pandas_frame["avocado_id"].astype("int32")
+    pandas_frame["avocado_id_int"] = pandas_frame["avocado_id"].astype("int32")
+    pandas_frame["avocado_id_int64"] = pandas_frame["avocado_id"].astype("int64")
+    pandas_frame["avocado_id_bigint"] = pandas_frame["avocado_id"].astype("int64")
+    pandas_frame["avocado_id_float"] = pandas_frame["avocado_id"].astype("float")
+    pandas_frame["avocado_id_float16"] = pandas_frame["avocado_id"].astype("float16")
+    pandas_frame["avocado_id_float32"] = pandas_frame["avocado_id"].astype("float32")
+    pandas_frame["avocado_id_float64"] = pandas_frame["avocado_id"].astype("float64")
+    pandas_frame["avocado_id_bool"] = pandas_frame["avocado_id"].astype("bool")
+    pandas_frame["avocado_id_category"] = pandas_frame["avocado_id"].astype("category")
+    pandas_frame["date"] = pandas_frame["Date"].astype("datetime64")
+    pandas_frame["time"] = pandas_frame["Date"].astype("datetime64")
+    pandas_frame["region_varchar"] = pandas_frame["region"].astype("string")
+    pandas_frame = pandas_frame.drop(columns=["avocado_id", "Date", "region"])
+
+    tm.assert_frame_equal(pandas_frame, my_frame)
+
+
 if __name__ == "__main__":
     register_env_tables()
 
-    test_in_operator()
+    test_sql_data_types()
 
     remove_env_tables()

dataframe_sql/tests/utils.py

@@ -1,21 +1,21 @@
 """
 Shared functions among the tests like setting up test environment
 """
-import os
 from pathlib import Path
 
 from pandas import DataFrame, read_csv
 
 from dataframe_sql import register_temp_table, remove_temp_table
 
-DATA_PATH = os.path.join(Path(__file__).parent.parent, "data")
+DATA_PATH = Path(__file__).parent.parent / "data"
 
 
 # Import the data for testing
-FOREST_FIRES = read_csv(os.path.join(DATA_PATH, "forestfires.csv"))
-DIGIMON_MON_LIST = read_csv(os.path.join(DATA_PATH, "DigiDB_digimonlist.csv"))
-DIGIMON_MOVE_LIST = read_csv(os.path.join(DATA_PATH, "DigiDB_movelist.csv"))
-DIGIMON_SUPPORT_LIST = read_csv(os.path.join(DATA_PATH, "DigiDB_supportlist.csv"))
+FOREST_FIRES = read_csv(DATA_PATH / "forestfires.csv")
+DIGIMON_MON_LIST = read_csv(DATA_PATH / "DigiDB_digimonlist.csv")
+DIGIMON_MOVE_LIST = read_csv(DATA_PATH / "DigiDB_movelist.csv")
+DIGIMON_SUPPORT_LIST = read_csv(DATA_PATH / "DigiDB_supportlist.csv")
+AVOCADO = read_csv(DATA_PATH / "avocado.csv")
 
 # Name change is for name interference
 DIGIMON_MON_LIST["mon_attribute"] = DIGIMON_MON_LIST["Attribute"]