Browse Source

[Filebrowser] Fix support for reading Parquet files (#3885)

Ayush Goyal 1 year ago
parent
commit
7b08df4b8b

+ 19 - 17
apps/filebrowser/src/filebrowser/views.py

@@ -32,9 +32,10 @@ from bz2 import decompress
 from datetime import datetime
 from datetime import datetime
 from functools import partial
 from functools import partial
 from gzip import decompress as decompress_gzip
 from gzip import decompress as decompress_gzip
-from io import StringIO as string_io
+from io import BytesIO, StringIO as string_io
 from urllib.parse import quote as urllib_quote, unquote as urllib_unquote, urlparse as lib_urlparse
 from urllib.parse import quote as urllib_quote, unquote as urllib_unquote, urlparse as lib_urlparse
 
 
+import pandas as pd
 from avro import datafile, io
 from avro import datafile, io
 from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload
 from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload
 from django.core.paginator import EmptyPage, InvalidPage, Page, Paginator
 from django.core.paginator import EmptyPage, InvalidPage, Page, Paginator
@@ -117,6 +118,9 @@ BYTES_PER_SENTENCE = 2
 # The maximum size the file editor will allow you to edit
 # The maximum size the file editor will allow you to edit
 MAX_FILEEDITOR_SIZE = 256 * 1024
 MAX_FILEEDITOR_SIZE = 256 * 1024
 
 
+# Parquet files start with a specific 4-byte magic number: 'PAR1'
+PARQUET_MAGIC_NUMBER = b'PAR1'
+
 INLINE_DISPLAY_MIMETYPE = re.compile(
 INLINE_DISPLAY_MIMETYPE = re.compile(
     r'video/|image/|audio/|application/pdf|application/msword|application/excel|application/vnd\.ms|application/vnd\.openxmlformats'
     r'video/|image/|audio/|application/pdf|application/msword|application/excel|application/vnd\.ms|application/vnd\.openxmlformats'
 )
 )
@@ -146,14 +150,6 @@ if hasattr(ARCHIVE_UPLOAD_TEMPDIR, 'get') and not os.path.exists(ARCHIVE_UPLOAD_
 logger = logging.getLogger()
 logger = logging.getLogger()
 
 
 
 
-class ParquetOptions(object):
-  def __init__(self, col=None, format='json', no_headers=True, limit=-1):
-    self.col = col
-    self.format = format
-    self.no_headers = no_headers
-    self.limit = limit
-
-
 def index(request):
 def index(request):
   # Redirect to home directory by default
   # Redirect to home directory by default
   path = request.user.get_home_directory()
   path = request.user.get_home_directory()
@@ -983,13 +979,16 @@ def _read_avro(fhandle, path, offset, length, stats):
 
 
 def _read_parquet(fhandle, path, offset, length, stats):
 def _read_parquet(fhandle, path, offset, length, stats):
   try:
   try:
-    size = 1 * 128 * 1024 * 1024  # Buffer file stream to 128 MB chunks
-    data = string_io(fhandle.read(size))
+    size = 1 * 128 * 1024 * 1024  # Buffer file stream to 128 MiB chunks
+
+    fhandle.seek(offset)
+    file_data = BytesIO(fhandle.read(size))
 
 
-    dumped_data = string_io()
-    parquet._dump(data, ParquetOptions(limit=1000), out=dumped_data)
-    dumped_data.seek(offset)
-    return dumped_data.read()
+    data_frame = pd.read_parquet(file_data, engine='pyarrow')
+
+    data_chunk = data_frame.iloc[offset:offset + length].to_string()
+
+    return data_chunk
   except Exception as e:
   except Exception as e:
     logging.exception('Could not read parquet file at "%s": %s' % (path, e))
     logging.exception('Could not read parquet file at "%s": %s' % (path, e))
     raise PopupException(_("Failed to read Parquet file."))
     raise PopupException(_("Failed to read Parquet file."))
@@ -1061,9 +1060,12 @@ def detect_snappy(contents):
 def detect_parquet(fhandle):
 def detect_parquet(fhandle):
   """
   """
   Detect parquet from magic header bytes.
   Detect parquet from magic header bytes.
-  Python 2 only currently.
   """
   """
-  return False if sys.version_info[0] > 2 else parquet._check_header_magic_bytes(fhandle)
+
+  fhandle.seek(0)
+  magic_number = fhandle.read(4)
+
+  return magic_number == PARQUET_MAGIC_NUMBER
 
 
 
 
 def snappy_installed():
 def snappy_installed():

+ 42 - 1
apps/filebrowser/src/filebrowser/views_test.py

@@ -24,11 +24,15 @@ import tempfile
 import urllib.error
 import urllib.error
 import urllib.parse
 import urllib.parse
 import urllib.request
 import urllib.request
+from io import BytesIO
 from time import sleep, time
 from time import sleep, time
 from unittest.mock import Mock, patch
 from unittest.mock import Mock, patch
 from urllib.parse import unquote as urllib_unquote
 from urllib.parse import unquote as urllib_unquote
 
 
+import pandas as pd
 import pytest
 import pytest
+import pyarrow as pa
+import pyarrow.parquet as pq
 from avro import datafile, io, schema
 from avro import datafile, io, schema
 from django.http import HttpResponse
 from django.http import HttpResponse
 from django.test import TestCase
 from django.test import TestCase
@@ -45,7 +49,7 @@ from desktop.lib.test_utils import add_permission, add_to_group, grant_access, r
 from desktop.lib.view_util import location_to_url
 from desktop.lib.view_util import location_to_url
 from filebrowser.conf import ENABLE_EXTRACT_UPLOADED_ARCHIVE, MAX_SNAPPY_DECOMPRESSION_SIZE, REMOTE_STORAGE_HOME
 from filebrowser.conf import ENABLE_EXTRACT_UPLOADED_ARCHIVE, MAX_SNAPPY_DECOMPRESSION_SIZE, REMOTE_STORAGE_HOME
 from filebrowser.lib.rwx import expand_mode
 from filebrowser.lib.rwx import expand_mode
-from filebrowser.views import _normalize_path, snappy_installed
+from filebrowser.views import _normalize_path, _read_parquet, snappy_installed
 from hadoop import pseudo_hdfs4
 from hadoop import pseudo_hdfs4
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.fs.webhdfs import WebHdfs
 from hadoop.fs.webhdfs import WebHdfs
@@ -1825,3 +1829,40 @@ class TestNormalizePath(object):
 
 
     normalized = _normalize_path(path)
     normalized = _normalize_path(path)
     assert path == normalized
     assert path == normalized
+
+
+class TestReadParquet:
+  def setup_method(self):
+    # Setup a common DataFrame and create a Parquet file in memory
+    self.test_df = pd.DataFrame({
+        'column1': [1, 2, 3, 4, 5],
+        'column2': ['a', 'b', 'c', 'd', 'e']
+    })
+    self.file_data = self.create_parquet_file(self.test_df)
+    self.path = "/mock/path/to/file.parquet"
+    self.offset = 0
+    self.length = 3
+    self.stats = None
+
+  def create_parquet_file(self, dataframe):
+    # Helper method to create a Parquet file in memory
+    buffer = BytesIO()
+    table = pa.Table.from_pandas(dataframe)
+    pq.write_table(table, buffer)
+    buffer.seek(0)  # Reset the file handle position
+    return buffer
+
+  def test_read_parquet_success(self):
+    # Call the function with valid Parquet data
+    result = _read_parquet(self.file_data, self.path, self.offset, self.length, self.stats)
+
+    expected_chunk = self.test_df.iloc[self.offset:self.offset + self.length].to_string()
+
+    assert result == expected_chunk
+
+  def test_read_parquet_invalid_file(self):
+    # Create an invalid file (not a Parquet file)
+    invalid_file_data = BytesIO(b"Not a valid Parquet file")
+
+    with pytest.raises(Exception, match="Failed to read Parquet file"):
+      _read_parquet(invalid_file_data, self.path, self.offset, self.length, self.stats)

+ 1 - 0
desktop/core/base_requirements.txt

@@ -45,6 +45,7 @@ phoenixdb==1.2.1
 prompt-toolkit==3.0.39
 prompt-toolkit==3.0.39
 protobuf==3.20.3
 protobuf==3.20.3
 py==1.11.0
 py==1.11.0
+pyarrow==17.0.0
 pyformance==0.3.2
 pyformance==0.3.2
 pytest==8.1.1
 pytest==8.1.1
 pytest-django==4.8.0
 pytest-django==4.8.0