# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. from random import choice, randint from http.server import BaseHTTPRequestHandler, HTTPServer import io import pyarrow as pa import pyarrow.compute as pc import re import socketserver import string # use dictionary encoding for the ticker column USE_DICTIONARY_ENCODING = True def random_string(alphabet, length): return "".join(choice(alphabet) for _ in range(length)) def random_name(initial): length = randint(3, 7) return initial + random_string(string.ascii_lowercase, length) def example_tickers(num_tickers): tickers = [] while len(tickers) < num_tickers: length = randint(3, 4) random_ticker = random_string(string.ascii_uppercase, length) if random_ticker not in tickers: tickers.append(random_ticker) return tickers the_ticker_type = ( pa.dictionary(pa.int32(), pa.utf8()) if USE_DICTIONARY_ENCODING else pa.utf8() ) the_schema = pa.schema( [ ("ticker", the_ticker_type), ("price", pa.int64()), ("volume", pa.int64()), ] ) def example_batch(tickers, length): ticker_indices = [] price = [] volume = [] for _ in range(length): ticker_indices.append(randint(0, len(tickers) - 1)) price.append(randint(1, 1000) * 100) volume.append(randint(1, 10000)) ticker = ( pa.DictionaryArray.from_arrays(ticker_indices, tickers) if USE_DICTIONARY_ENCODING else pc.take(tickers, ticker_indices, boundscheck=False) ) return pa.RecordBatch.from_arrays([ticker, price, volume], schema=the_schema) def example_batches(tickers): # these parameters are chosen to generate a response # of ~1 GB and chunks of ~140 KB (uncompressed) total_records = 42_000_000 batch_len = 6 * 1024 # all the batches sent are random slices of the larger base batch base_batch = example_batch(tickers, length=8 * batch_len) batches = [] records = 0 while records < total_records: length = min(batch_len, total_records - records) offset = randint(0, base_batch.num_rows - length - 1) batch = base_batch.slice(offset, length) batches.append(batch) records += length return batches # end of example data generation # what the HTTP spec calls a token (any character except CTLs or separators) TOKEN_RE = r"(?:[A-Za-z0-9!#$%&'*+./^_`|~-]+)" # [L]inear [W]hite [S]pace pattern (HTTP/1.1 - RFC 2616) LWS_RE = r"(?:[ \t]|\r\n[ \t]+)*" TOKENIZER_PAT = re.compile( f"(?P{TOKEN_RE})" r'|(?P"([^"\\]|\\.)*")' # a quoted string (escaped pairs allowed) r"|(?P,)" r"|(?P;)" r"|(?P=)" f"|(?P{LWS_RE})" # LWS is skipped r"|(?P.+)", flags=re.ASCII, # HTTP headers are encoded in ASCII ) def parse_header_value(header_name, header_value): """ Parse the Accept or Accept-Encoding request header values. Returns ------- list of (str, dict) The list of lowercase tokens and their parameters in the order they appear in the header. The parameters are stored in a dictionary where the keys are the parameter names and the values are the parameter values. If a parameter is not followed by an equal sign and a value, the value is None. """ def unexpected(label, value): msg = f"Malformed {header_name} header: unexpected {label} at {value!r}" return ValueError(msg) def tokenize(): for mo in re.finditer(TOKENIZER_PAT, header_value): kind = mo.lastgroup if kind == "SKIP": continue elif kind == "MISMATCH": raise unexpected("character", mo.group()) yield (kind, mo.group()) tokens = tokenize() def expect(expected_kind): kind, text = next(tokens) if kind != expected_kind: raise unexpected("token", text) return text accepted = [] while True: try: name, params = None, {} name = expect("TOK").lower() kind, text = next(tokens) while True: if kind == "COMMA": accepted.append((name, params)) break if kind == "SEMI": ident = expect("TOK") params[ident] = None # init param to None kind, text = next(tokens) if kind != "EQ": continue kind, text = next(tokens) if kind in ["TOK", "QUOTED"]: if kind == "QUOTED": text = text[1:-1] # remove the quotes params[ident] = text # set param to value kind, text = next(tokens) continue raise unexpected("token", text) except StopIteration: break if name is not None: # any unfinished ;param=value sequence or trailing separators are ignored accepted.append((name, params)) return accepted ARROW_STREAM_FORMAT = "application/vnd.apache.arrow.stream" def pick_ipc_codec(accept_header, available, default): """ Pick the IPC stream codec according to the Accept header. This is used when deciding which codec to use for compression of IPC buffer streams. This is a feature of the Arrow IPC stream format and is different from the HTTP content-coding used to compress the entire HTTP response. This is how a client may specify the IPC buffer compression codecs it accepts: Accept: application/vnd.apache.arrow.stream; codecs="zstd, lz4" Parameters ---------- accept_header : str|None The value of the Accept header from an HTTP request. available : list of str The codecs that the server can provide in the order preferred by the server. Example: ["zstd", "lz4"]. default : str|None The codec to use if the client does not specify the ";codecs" parameter in the Accept header. Returns ------- str|None The codec that the server should use to compress the IPC buffer stream. None if the client does not accept any of the available codecs explicitly listed. ;codecs="" means no codecs are accepted. If the client does not specify the codecs parameter, the default codec is returned. """ did_specify_codecs = False accepted_codecs = [] if accept_header is not None: accepted = parse_header_value("Accept", accept_header) for media_range, params in accepted: if ( media_range == "*/*" or media_range == "application/*" or media_range == ARROW_STREAM_FORMAT ): did_specify_codecs = "codecs" in params codecs_str = params.get("codecs") if codecs_str is None: continue for codec in codecs_str.split(","): accepted_codecs.append(codec.strip()) for codec in available: if codec in accepted_codecs: return codec return None if did_specify_codecs else default def pick_coding(accept_encoding_header, available): """ Pick the content-coding according to the Accept-Encoding header. This is used when using HTTP response compression instead of IPC buffer compression. Parameters ---------- accept_encoding_header : str The value of the Accept-Encoding header from an HTTP request. available : list of str The content-codings that the server can provide in the order preferred by the server. Example: ["zstd", "br", "gzip"]. Returns ------- str The content-coding that the server should use to compress the response. "identity" is returned if no acceptable content-coding is found in the list of available codings. None if the client does not accept any of the available content-codings and doesn't accept "identity" (uncompressed) either. In this case, a "406 Not Acceptable" response should be sent. """ accepted = parse_header_value("Accept-Encoding", accept_encoding_header) def qvalue_or(params, default): qvalue = params.get("q") if qvalue is not None: try: return float(qvalue) except ValueError: raise ValueError(f"Invalid qvalue in Accept-Encoding header: {qvalue}") return default if "identity" not in available: available = available + ["identity"] state = {} for coding, params in accepted: qvalue = qvalue_or(params, 0.0001 if coding == "identity" else 1.0) if coding == "*": for coding in available: if coding not in state: state[coding] = qvalue elif coding in available: state[coding] = qvalue # "identity" is always acceptable unless explicitly refused (;q=0) if "identity" not in state: state["identity"] = 0.0001 # all the candidate codings are now in the state dictionary and we # have to consider only the ones that have the maximum qvalue max_qvalue = max(state.values()) if max_qvalue == 0.0: return None for coding in available: if coding in state and state[coding] == max_qvalue: return coding return None def pick_compression(headers, available_ipc_codecs, available_codings, default): """ Pick the compression strategy based on the Accept and Accept-Encoding headers. Parameters ---------- headers : dict The HTTP request headers. available_ipc_codecs : list of str The codecs that the server can provide for IPC buffer compression. available_codings : list of str The content-codings that the server can provide for HTTP response compression. default : str The default compression strategy to use if the client does explicitly choose. Returns ------- str|None The compression strategy to use. It can be one of the following: "identity": no compression at all. "identity+zstd": No HTTP compression + IPC buffer compression with Zstd. "identity+lz4": No HTTP compression + IPC buffer compression with LZ4. "zstd", "br", "gzip", ...: HTTP compression without IPC buffer compression. None means a "406 Not Acceptable" response should be sent. """ accept = headers.get("Accept") ipc_codec = pick_ipc_codec(accept, available_ipc_codecs, default=None) if ipc_codec is None: accept_encoding = headers.get("Accept-Encoding") return ( default if accept_encoding is None else pick_coding(accept_encoding, available_codings) ) return "identity+" + ipc_codec class LateClosingBytesIO(io.BytesIO): """ BytesIO that does not close on close(). When a stream wrapping a a file-like object is closed, the underlying file-like object is also closed. This subclass prevents that from happening by overriding the close method. If we close a RecordBatchStreamWriter wrapping a BytesIO object, we want to be able to create a memory view of the buffer. But that is only possible if the BytesIO object is not closed yet. """ def close(self): pass def close_now(self): super().close() class SocketWriterSink(socketserver._SocketWriter): """Wrapper to make wfile usable as a sink for Arrow stream writing.""" def __init__(self, wfile): self.wfile = wfile def writable(self): return True def write(self, b): self.wfile.write(b) def fileno(self): return self._sock.fileno() def close(self): """Do nothing so Arrow stream wrappers don't close the socket.""" pass def generate_chunk_buffers(schema, source, compression): # the sink holds the buffer and we give a view of it to the caller with LateClosingBytesIO() as sink: # keep buffering until we have at least MIN_BUFFER_SIZE bytes # in the buffer before yielding it to the caller. Setting it # to 1 means we yield as soon as the compression blocks are # formed and reach the sink buffer. MIN_BUFFER_SIZE = 64 * 1024 if compression.startswith("identity"): if compression == "identity+zstd": options = pa.ipc.IpcWriteOptions(compression="zstd") elif compression == "identity+lz4": options = pa.ipc.IpcWriteOptions(compression="lz4") else: options = None # source: RecordBatchReader # |> writer: RecordBatchStreamWriter # |> sink: LateClosingBytesIO writer = pa.ipc.new_stream(sink, schema, options=options) for batch in source: writer.write_batch(batch) if sink.tell() >= MIN_BUFFER_SIZE: sink.truncate() with sink.getbuffer() as buffer: yield buffer sink.seek(0) writer.close() # write EOS marker and flush else: compression = "brotli" if compression == "br" else compression with pa.CompressedOutputStream(sink, compression) as compressor: # has the first buffer been yielded already? sent_first = False # source: RecordBatchReader # |> writer: RecordBatchStreamWriter # |> compressor: CompressedOutputStream # |> sink: LateClosingBytesIO writer = pa.ipc.new_stream(compressor, schema) for batch in source: writer.write_batch(batch) # we try to yield a buffer ASAP no matter how small if not sent_first and sink.tell() == 0: compressor.flush() pos = sink.tell() if pos >= MIN_BUFFER_SIZE or (not sent_first and pos >= 1): sink.truncate() with sink.getbuffer() as buffer: yield buffer sink.seek(0) sent_first = True writer.close() # write EOS marker and flush compressor.close() sink.truncate() with sink.getbuffer() as buffer: yield buffer sink.close_now() AVAILABLE_IPC_CODECS = ["zstd", "lz4"] """List of available codecs Arrow IPC buffer compression.""" AVAILABLE_CODINGS = ["zstd", "br", "gzip"] """ List of available content-codings as used in HTTP. Note that Arrow stream classes refer to Brotli as "brotli" and not "br". """ class MyRequestHandler(BaseHTTPRequestHandler): """ Response handler for a simple HTTP server. This HTTP request handler serves a compressed HTTP response with an Arrow stream in it or a (TODO) compressed Arrow stream in a uncompressed HTTP response. The Arrow data is randomly generated "trading data" with a schema consisting of a ticker, price (in cents), and volume. """ def _resolve_batches(self): return pa.RecordBatchReader.from_batches(the_schema, all_batches) def _send_not_acceptable(self, parsing_error=None): self.send_response(406, "Not Acceptable") self.send_header("Content-Type", "text/plain") self.end_headers() if parsing_error: message = f"Error parsing header: {parsing_error}\n" else: message = "None of the available codings are accepted by this client.\n" accept = self.headers.get("Accept") if accept is not None: message += f"`Accept` header was {accept!r}.\n" accept_encoding = self.headers.get("Accept-Encoding") if accept_encoding is not None: message += f"`Accept-Encoding` header was {accept_encoding!r}.\n" self.wfile.write(bytes(message, "utf-8")) def do_GET(self): # HTTP/1.0 requests don't get chunked responses if self.request_version == "HTTP/1.0": self.protocol_version = "HTTP/1.0" chunked = False else: self.protocol_version = "HTTP/1.1" chunked = True # if client's intent cannot be derived from the headers, return # uncompressed data for HTTP/1.0 requests and compressed data for # HTTP/1.1 requests with the safest compression format choice: "gzip". default_compression = ( "identity" if self.request_version == "HTTP/1.0" or ("gzip" not in AVAILABLE_CODINGS) else "gzip" ) try: compression = pick_compression( self.headers, AVAILABLE_IPC_CODECS, AVAILABLE_CODINGS, default_compression, ) if compression is None: self._send_not_acceptable() return except ValueError as e: self._send_not_acceptable(str(e)) return ### in a real application the data would be resolved from a database or ### another source like a file and error handling would be done here ### before the 200 OK response starts being sent to the client. source = self._resolve_batches() self.send_response(200) ### set these headers if testing with a local browser-based client: # self.send_header('Access-Control-Allow-Origin', 'http://localhost:8008') # self.send_header('Access-Control-Allow-Methods', 'GET') # self.send_header('Access-Control-Allow-Headers', 'Content-Type') self.send_header( "Content-Type", ( f"{ARROW_STREAM_FORMAT}; codecs={compression[9:]}" if compression.startswith("identity+") else ARROW_STREAM_FORMAT ), ) # suggest a default filename in case this response is saved by the user self.send_header("Content-Disposition", r'attachment; filename="output.arrows"') if not compression.startswith("identity"): self.send_header("Content-Encoding", compression) if chunked: self.send_header("Transfer-Encoding", "chunked") self.end_headers() for buffer in generate_chunk_buffers(the_schema, source, compression): self.wfile.write(f"{len(buffer):X}\r\n".encode("utf-8")) self.wfile.write(buffer) self.wfile.write("\r\n".encode("utf-8")) self.wfile.write("0\r\n\r\n".encode("utf-8")) else: self.end_headers() sink = SocketWriterSink(self.wfile) for buffer in generate_chunk_buffers(the_schema, source, compression): sink.write(buffer) print("Generating example data...") all_tickers = example_tickers(60) all_batches = example_batches(all_tickers) server_address = ("localhost", 8008) try: httpd = HTTPServer(server_address, MyRequestHandler) print(f"Serving on {server_address[0]}:{server_address[1]}...") httpd.serve_forever() except KeyboardInterrupt: print("Shutting down server") httpd.socket.close()