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 functools import partial
 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
 
+import pandas as pd
 from avro import datafile, io
 from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload
 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
 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(
     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()
 
 
-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):
   # Redirect to home directory by default
   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):
   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:
     logging.exception('Could not read parquet file at "%s": %s' % (path, e))
     raise PopupException(_("Failed to read Parquet file."))
@@ -1061,9 +1060,12 @@ def detect_snappy(contents):
 def detect_parquet(fhandle):
   """
   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():

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

@@ -24,11 +24,15 @@ import tempfile
 import urllib.error
 import urllib.parse
 import urllib.request
+from io import BytesIO
 from time import sleep, time
 from unittest.mock import Mock, patch
 from urllib.parse import unquote as urllib_unquote
 
+import pandas as pd
 import pytest
+import pyarrow as pa
+import pyarrow.parquet as pq
 from avro import datafile, io, schema
 from django.http import HttpResponse
 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 filebrowser.conf import ENABLE_EXTRACT_UPLOADED_ARCHIVE, MAX_SNAPPY_DECOMPRESSION_SIZE, REMOTE_STORAGE_HOME
 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.conf import UPLOAD_CHUNK_SIZE
 from hadoop.fs.webhdfs import WebHdfs
@@ -1825,3 +1829,40 @@ class TestNormalizePath(object):
 
     normalized = _normalize_path(path)
     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
 protobuf==3.20.3
 py==1.11.0
+pyarrow==17.0.0
 pyformance==0.3.2
 pytest==8.1.1
 pytest-django==4.8.0