teradataml 20.0.0.5__py3-none-any.whl → 20.0.0.7__py3-none-any.whl

teradataml/dataframe/dataframe_utils.py

@@ -1796,7 +1796,8 @@ class DataFrameUtils():
         fil_nodeid = aed_utils._aed_filter(sel_nodeid, filter_str)
         sel2_nodeid = aed_utils._aed_select(fil_nodeid, sel_cols_str)
         col_names, col_types = __class__._get_column_names_and_types_from_metaexpr(df._metaexpr)
-        new_metaexpr = UtilFuncs._get_metaexpr_using_columns(df._nodeid, zip(col_names, col_types))
+        new_metaexpr = UtilFuncs._get_metaexpr_using_columns(df._nodeid, zip(col_names, col_types),
+                                                             datalake=df._metaexpr.datalake)
         # Call the function from_node from appropriate class either DataFrame or GeoDataFrame
         new_df = df.__class__._from_node(sel2_nodeid, new_metaexpr, df._index_label)
         new_df._orderby = df._orderby
@@ -1855,6 +1856,8 @@ class DataFrameUtils():
 
             db_schema = UtilFuncs._extract_db_name(tab_name_first)
             db_table_name = UtilFuncs._extract_table_name(tab_name_first)
+            if dfs[0]._metaexpr.datalake:
+                return DataFrame(in_schema(db_schema, db_table_name, dfs[0]._metaexpr.datalake))
 
             if db_schema:
                 return DataFrame(in_schema(db_schema, db_table_name))
@@ -1875,7 +1878,9 @@ class DataFrameUtils():
                 db_schema = UtilFuncs._extract_db_name(tab_name_first)
                 db_table_name = UtilFuncs._extract_table_name(tab_name_first)
 
-                if db_schema:
+                if dfs[i]._metaexpr.datalake:
+                    parent_df = DataFrame(in_schema(db_schema, db_table_name, dfs[i]._metaexpr.datalake))
+                elif db_schema:
                     parent_df = DataFrame(in_schema(db_schema, db_table_name))
                 else:
                     parent_df = DataFrame(db_table_name)
@@ -1919,7 +1924,7 @@ class DataFrameUtils():
         # 2. Comma separated parameters enclosed in parentheses
         # 3. Comma separated parameters without parenthesis 
         # 4. Remaining string
-        pattern = "([A-Z0-9_]+)(\((.*)\))?(.*)"
+        pattern = r"([A-Z0-9_]+)(\((.*)\))?(.*)"
 
         m = re.match(pattern, td_type)
         td_str_type = m.group(1)
@@ -1953,7 +1958,7 @@ class DataFrameUtils():
         return _DtypesMappers.DATALAKE_STR_to_TDSQLALCHEMY_DATATYPE_MAPPER[td_str_type]()
 
     @staticmethod
-    def _get_datalake_table_columns_info(schema, table_name, datalake):
+    def _get_datalake_table_columns_info(schema, table_name, datalake, use_dialect=False):
         """
         Function to get column names and corresponding teradatasqlalchemy types
         of a datalake table using results of 'help table <datalake>.<db_name>.<table_name>'
@@ -1990,31 +1995,64 @@ class DataFrameUtils():
               VARCHAR(length=2000, charset='UNICODE'),
               INTEGER()])
         """
-        # Get the column information from the strings type.
-        prepared = preparer(td_dialect())
-        sqlbundle = SQLBundle()
-        full_tbl_name = '{}.{}.{}'.format(prepared.quote(datalake),
-                                          prepared.quote(schema),
-                                          prepared.quote(table_name))
-        help_table_sql = sqlbundle._get_sql_query(SQLConstants.SQL_HELP_TABLE).format(full_tbl_name)
-
-        cur = execute_sql(help_table_sql)
-        td_types_col_index = -1
-        for i, col_metadata in enumerate(cur.description):
-            # Help Table returns column names and
-            # corresponding IcebergType, TeradataInternalType,
-            # TeradataType. We need to extract column index for
-            # 'TeradataType' column.
-            if col_metadata[0].lower() == 'teradatatype':
-                td_types_col_index = i
-
         col_names = []
         col_types = []
-        if td_types_col_index > -1:
-            for col_info in cur.fetchall():
-                col_names.append(col_info[0])
-                col_types.append(DataFrameUtils._get_sqlalchemy_type_from_str(col_info[td_types_col_index]))
+        if not use_dialect:
+            # Get the column information from the strings type.
+            prepared = preparer(td_dialect())
+            sqlbundle = SQLBundle()
+            full_tbl_name = '{}.{}.{}'.format(prepared.quote(datalake),
+                                              prepared.quote(schema),
+                                              prepared.quote(table_name))
+            help_table_sql = sqlbundle._get_sql_query(SQLConstants.SQL_HELP_TABLE).format(full_tbl_name)
+
+            cur = execute_sql(help_table_sql)
+            td_types_col_index = -1
+
+            for i, col_metadata in enumerate(cur.description):
+                # Help Table returns column names and
+                # corresponding IcebergType, TeradataInternalType,
+                # TeradataType. We need to extract column index for
+                # 'TeradataType' column.
+                if col_metadata[0].lower() in ['teradatatype', 'Type']:
+                    td_types_col_index = i
+
+            if td_types_col_index > -1:
+                for col_info in cur.fetchall():
+                    col_names.append(col_info[0])
+                    col_types.append(DataFrameUtils._get_sqlalchemy_type_from_str(col_info[td_types_col_index]))
+            else:
+                raise TeradataMlException(Messages.get_message(MessageCodes.TDMLDF_CREATE_FAIL),
+                                          MessageCodes.TDMLDF_CREATE_FAIL)
         else:
-            raise TeradataMlException(Messages.get_message(MessageCodes.TDMLDF_CREATE_FAIL),
-                                      MessageCodes.TDMLDF_CREATE_FAIL)
+            new_kwarg = get_connection().dialect.__class__.__name__ + "_datalake"
+            all_col_info = get_connection().dialect.get_columns(connection=get_connection(),
+                                                                table_name=table_name,
+                                                                schema=schema,
+                                                                table_only=True,
+                                                                **{new_kwarg: datalake})
+            for col_dict in all_col_info:
+                col_names.append(col_dict.get('name', col_dict.get('Column Name')))
+                col_types.append(col_dict.get('type', col_dict.get('Type')))
+
         return col_names, col_types
+
+    @staticmethod
+    def check_otf_dataframe():
+        """Decorator for validating if DataFrame is created on OTF table or not and throw error."""
+        def decorator(method):
+            def wrapper(self, *args, **kwargs):
+                if not self._datalake:
+                    attr = getattr(type(self), method.__name__, None)
+                    caller_name = method.__name__ + '()'
+                    if isinstance(attr, property):
+                        caller_name = method.__name__
+                    raise TeradataMlException(Messages.get_message(MessageCodes.OTF_TABLE_REQUIRED,
+                                                                   caller_name),
+                                              MessageCodes.UNSUPPORTED_OPERATION)
+
+                return method(self, *args, **kwargs)
+
+            return wrapper
+
+        return decorator