فهرست منبع

HUE-4501 [indexer] Add basic hive CSV file format outline to indexer

peddle 9 سال پیش
والد
کامیت
758c3fc
1فایلهای تغییر یافته به همراه47 افزوده شده و 6 حذف شده
  1. 47 6
      desktop/libs/indexer/src/indexer/file_format.py

+ 47 - 6
desktop/libs/indexer/src/indexer/file_format.py

@@ -287,9 +287,7 @@ class CSVFormat(FileFormat):
 
   @classmethod
   def _guess_from_file_stream(cls, file_stream):
-    file_stream.seek(0)
-    sample = '\n'.join(file_stream.read(1024*1024*5).splitlines())
-    file_stream.seek(0)
+    sample = cls._get_sample(file_stream)
 
     try:
       dialect, has_header = cls._guess_dialect(sample)
@@ -309,11 +307,17 @@ class CSVFormat(FileFormat):
     })
 
   @classmethod
-  def _from_format(cls, file_stream, format_):
+  def _get_sample(cls, file_stream):
     file_stream.seek(0)
     sample = '\n'.join(file_stream.read(1024*1024*5).splitlines())
     file_stream.seek(0)
 
+    return sample
+
+  @classmethod
+  def _from_format(cls, file_stream, format_):
+    sample = cls._get_sample(file_stream)
+
     delimiter = format_["fieldSeparator"].encode('utf-8')
     line_terminator = format_["recordSeparator"].encode('utf-8')
     quote_char = format_["quoteChar"].encode('utf-8')
@@ -333,7 +337,7 @@ class CSVFormat(FileFormat):
     else:
       return cls._guess_from_file_stream(file_stream)
 
-  def __init__(self, delimiter=',', line_terminator='\n', quote_char='"', has_header=False, sample=""):
+  def __init__(self, delimiter=',', line_terminator='\n', quote_char='"', has_header=False, sample="", fields=None):
     self._delimiter = delimiter
     self._line_terminator = line_terminator
     self._quote_char = quote_char
@@ -346,7 +350,7 @@ class CSVFormat(FileFormat):
 
     self._num_columns = self._guess_num_columns(self._sample_rows)
 
-    self._fields = self._guess_fields(sample)
+    self._fields = fields if fields else self._guess_fields(sample)
 
     super(CSVFormat, self).__init__()
 
@@ -450,3 +454,40 @@ class CSVFormat(FileFormat):
       fields = []
 
     return fields
+
+class HiveFormat(CSVFormat):
+  FIELD_TYPE_TRANSLATE = {
+    "BOOLEAN_TYPE": "string",
+    "TINYINT_TYPE": "long",
+    "SMALLINT_TYPE": "long",
+    "INT_TYPE": "long",
+    "BIGINT_TYPE": "long",
+    "FLOAT_TYPE": "double",
+    "DOUBLE_TYPE": "double",
+    "STRING_TYPE":"string",
+    "TIMESTAMP_TYPE": "date",
+    "BINARY_TYPE": "long",
+    "DECIMAL_TYPE": "double",
+    "DATE_TYPE": "date",
+  }
+
+  @classmethod
+  def get_instance(cls, file_stream, format_):
+    sample = cls._get_sample(file_stream)
+
+    fields = []
+
+    for field in format_["fields"]:
+      fields.append(Field(
+        name=field["name"],
+        field_type_name=cls.FIELD_TYPE_TRANSLATE.get(field['type'], 'string')
+        ))
+
+    return cls(**{
+      "delimiter":',',
+      "line_terminator": '\n',
+      "quote_char": '"',
+      "has_header": False,
+      "sample": sample,
+      "fields": format_["fields"]
+    })