<commit_before>(ns flambo.parquet.thrift
  (:import [parquet.hadoop ParquetOutputFormat]
           [parquet.hadoop.thrift
            ParquetThriftOutputFormat
            ParquetThriftInputFormat
            ThriftReadSupport]
           [parquet.hadoop.metadata CompressionCodecName]
           [org.apache.hadoop.mapreduce Job]
           [org.apache.hadoop.io NullWritable]))

(def COMPRESSION-CODEC {:uncompressed CompressionCodecName/UNCOMPRESSED
                        :gzip CompressionCodecName/GZIP
                        :snappy CompressionCodecName/SNAPPY
                        :lzo CompressionCodecName/LZO
                        })

;; # InputFormat configuration
;;
(defn unbound-record-filter! [job filter-sym]
  (doto job
    (ParquetThriftInputFormat/setUnboundRecordFilter (Class/forName filter-sym))))

(defn column-filter! [conf filter]
  (doto conf
    (.set ThriftReadSupport/THRIFT_COLUMN_FILTER_KEY filter)))

(defn parquet-thrift-file
  "Create an RDD from a directory of parquet-thrift files
  where `klass` is the thrift class used to create the parquet files."
  [spark-context path klass & {:keys [job unbound-record-filter filter]
                               :or {job (Job.)}}]
  (let [job (if unbound-record-filter
              (unbound-record-filter! job unbound-record-filter)
              job)
        conf (if filter
               (-> (.getConfiguration job) (column-filter! filter))<commit_msg>add link to filter docs<commit_after>(ns flambo.parquet.thrift
  (:import [parquet.hadoop ParquetOutputFormat]
           [parquet.hadoop.thrift
            ParquetThriftOutputFormat
            ParquetThriftInputFormat
            ThriftReadSupport]
           [parquet.hadoop.metadata CompressionCodecName]
           [org.apache.hadoop.mapreduce Job]
           [org.apache.hadoop.io NullWritable]))

(def COMPRESSION-CODEC {:uncompressed CompressionCodecName/UNCOMPRESSED
                        :gzip CompressionCodecName/GZIP
                        :snappy CompressionCodecName/SNAPPY
                        :lzo CompressionCodecName/LZO
                        })

;; # InputFormat configuration
;;
(defn unbound-record-filter! [job filter-sym]
  (doto job
    (ParquetThriftInputFormat/setUnboundRecordFilter (Class/forName filter-sym))))

(defn column-filter! [conf filter]
  (doto conf
    (.set ThriftReadSupport/THRIFT_COLUMN_FILTER_KEY filter)))

(defn parquet-thrift-file
  "Create an RDD from a directory of parquet-thrift files
  where `klass` is the thrift class used to create the parquet files.

  See https://github.com/Parquet/parquet-mr/blob/master/parquet_cascading.md
  for the filter syntax."
  [spark-context path klass & {:keys [job unbound-record-filter filter]
                               :or {job (Job.)}}]
  (let [job (if unbound-record-filter
              (unbound-record-filter! job unbound-record-filter)
              job)
        conf (if filter
               (-> (.getConfiguration job) (column-filter! filter))
{'fraction_non_alphanumeric': 0.09009009009009009, 'fraction_numerical': 0.0, 'mean_word_length': 2.4106939704209327, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '900668', 'n_tokens_mistral': 867, 'n_tokens_neox': 866, 'n_words': 195}